├── .flake8 ├── .github └── workflows │ └── main.yml ├── .gitignore ├── .isort.cfg ├── .pylintrc ├── .style.yapf ├── CHANGELOG.md ├── LICENSE ├── README.md ├── README.rst ├── configs ├── classification │ ├── efficientnet_v1_b0.yml │ ├── efficientnet_v2_b0.yml │ ├── efficientnet_v2_small.yml │ ├── mobilenetv3_large.yml │ ├── mobilenetv3_large_075.yml │ ├── mobilenetv3_large_aux.yml │ └── mobilenetv3_small.yml ├── multilabel_classification │ ├── efficientnet_v1_b0.yml │ ├── efficientnet_v2_b0.yml │ ├── efficientnet_v2_small.yml │ ├── mobilenetv3_large.yml │ ├── mobilenetv3_large_75.yml │ └── mobilenetv3_small.yml └── opt_configs │ └── example.json ├── constraints.txt ├── init_venv.sh ├── linter.sh ├── openvino-requirements.txt ├── optional-requirements.txt ├── requirements.txt ├── scripts ├── default_config.py └── script_utils.py ├── setup.py ├── third-party-programs.txt ├── tools ├── auxiliary │ ├── add_initial_lr_to_checkpoint.py │ ├── add_nncf_metainfo_to_checkpoint.py │ ├── compute_mean_std.py │ ├── explore_weighs.py │ ├── optuna_optim.py │ ├── parse_output.sh │ ├── run_classification.py │ ├── run_multilabel.py │ └── visualize_actmap.py ├── check_classes_mlc.py ├── convert_to_onnx.py ├── eval.py ├── get_flops.py ├── main.py └── run_pylint.py └── torchreid ├── __init__.py ├── apis ├── export.py └── training.py ├── data ├── __init__.py ├── datamanager.py ├── datasets │ ├── __init__.py │ ├── dataset.py │ └── image │ │ ├── __init__.py │ │ └── classification.py ├── sampler.py └── transforms.py ├── engine ├── __init__.py ├── builder.py ├── engine.py └── image │ ├── __init__.py │ ├── am_softmax.py │ ├── multihead.py │ └── multilabel.py ├── integration ├── __init__.py └── nncf │ ├── __init__.py │ ├── accuracy_aware_training │ ├── __init__.py │ ├── base_runner.py │ ├── runner.py │ └── training_loop.py │ ├── compression.py │ ├── compression_script_utils.py │ ├── config.py │ └── engine.py ├── losses ├── __init__.py ├── am_softmax.py ├── asl.py └── cross_entropy_loss.py ├── metrics ├── __init__.py ├── accuracy.py └── classification.py ├── models ├── __init__.py ├── common.py ├── efficient_net_pytcv.py ├── gcn.py ├── inceptionv4_pytcv.py ├── ml_decoder.py ├── mobilenetv3.py ├── model_store.py ├── ptcv_wrapper.py ├── q2l.py ├── timm_wrapper.py └── transformer.py ├── ops ├── __init__.py ├── data_parallel.py ├── dropout.py └── self_challenging.py ├── optim ├── __init__.py ├── lr_finder.py ├── lr_scheduler.py ├── optimizer.py ├── radam.py └── sam.py ├── utils ├── __init__.py ├── avgmeter.py ├── fmix.py ├── ie_tools.py ├── loggers.py ├── tools.py └── torchtools.py └── version.py /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | ignore = 3 | # At least two spaces before inline comment 4 | E261, 5 | # Line lengths are recommended to be no greater than 79 characters 6 | E501, 7 | # Missing whitespace around arithmetic operator 8 | E226, 9 | # Blank line contains whitespace 10 | W293, 11 | # Do not use bare 'except' 12 | E722, 13 | # Line break after binary operator 14 | W504, 15 | # isort found an import in the wrong position 16 | I001 17 | max-line-length = 79 18 | exclude = __init__.py, build, torchreid/metrics/rank_cylib/ -------------------------------------------------------------------------------- /.github/workflows/main.yml: -------------------------------------------------------------------------------- 1 | name: Python package 2 | 3 | on: 4 | pull_request: 5 | push: 6 | branches: [main] 7 | 8 | jobs: 9 | build: 10 | 11 | runs-on: ubuntu-latest 12 | strategy: 13 | matrix: 14 | python-version: [3.8] 15 | 16 | steps: 17 | - name: Checkout repository and submodules 18 | uses: actions/checkout@v2 19 | with: 20 | submodules: recursive 21 | - name: Set up Python ${{ matrix.python-version }} 22 | uses: actions/setup-python@v2 23 | with: 24 | python-version: ${{ matrix.python-version }} 25 | - name: Cache pip 26 | uses: actions/cache@v2 27 | with: 28 | # This path is specific to Ubuntu 29 | path: ~/.cache/pip 30 | # Look to see if there is a cache hit for the corresponding requirements file 31 | key: ${{ runner.os }}-pip-${{ hashFiles('requirements.txt') }} 32 | restore-keys: | 33 | ${{ runner.os }}-pip- 34 | ${{ runner.os }}- 35 | - name: Install dependencies 36 | run: | 37 | python -m pip install --upgrade pip pytest 38 | if [ -f requirements.txt ]; then pip install -r requirements.txt; fi 39 | if [ -f optional-requirements.txt ]; then pip install -r optional-requirements.txt; fi 40 | if [ -f openvino-requirements.txt ]; then pip install -r openvino-requirements.txt; fi 41 | - name: Install torchreid 42 | run: | 43 | python setup.py develop 44 | - name: Linting with pylint 45 | run: | 46 | python tools/run_pylint.py 47 | -------------------------------------------------------------------------------- /.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 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | 62 | # Flask stuff: 63 | instance/ 64 | .webassets-cache 65 | 66 | # Scrapy stuff: 67 | .scrapy 68 | 69 | # Sphinx documentation 70 | docs/_build/ 71 | 72 | # PyBuilder 73 | target/ 74 | 75 | # Jupyter Notebook 76 | .ipynb_checkpoints 77 | 78 | # IPython 79 | profile_default/ 80 | ipython_config.py 81 | 82 | # pyenv 83 | .python-version 84 | 85 | # pipenv 86 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 87 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 88 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 89 | # install all needed dependencies. 90 | #Pipfile.lock 91 | 92 | # celery beat schedule file 93 | celerybeat-schedule 94 | 95 | # SageMath parsed files 96 | *.sage.py 97 | 98 | # Environments 99 | .env 100 | .venv 101 | env/ 102 | venv/ 103 | ENV/ 104 | env.bak/ 105 | venv.bak/ 106 | 107 | # Spyder project settings 108 | .spyderproject 109 | .spyproject 110 | 111 | # Rope project settings 112 | .ropeproject 113 | 114 | # mkdocs documentation 115 | /site 116 | 117 | # mypy 118 | .mypy_cache/ 119 | .dmypy.json 120 | dmypy.json 121 | 122 | # Pyre type checker 123 | .pyre/ 124 | 125 | # Cython eval code 126 | *.c 127 | *.html 128 | 129 | # OS X 130 | .DS_Store 131 | .Spotlight-V100 132 | .Trashes 133 | ._* 134 | 135 | # ReID 136 | reid-data/ 137 | log/ 138 | saved-models/ 139 | model-zoo/ 140 | debug* 141 | log* 142 | 143 | .history 144 | -------------------------------------------------------------------------------- /.isort.cfg: -------------------------------------------------------------------------------- 1 | [isort] 2 | line_length = 79 3 | multi_line_output = 0 4 | known_standard_library = setuptools 5 | known_first_party = torchreid 6 | known_third_party = PIL,cv2,numpy,Cython,h5py,yaml,Pillow,six,scipy,opencv-python,matplotlib,tb-nightly,future,yacs,gdown,flake8,yapf,terminaltables,scikit-learn,defusedxml,pytorchcv,terminaltables,torch_lr_finder,onnx,torch,torchvision 7 | no_lines_before = STDLIB,LOCALFOLDER 8 | default_section = THIRDPARTY -------------------------------------------------------------------------------- /.pylintrc: -------------------------------------------------------------------------------- 1 | [MASTER] 2 | disable = arguments-differ, 3 | cell-var-from-loop, 4 | fixme, 5 | global-statement, 6 | invalid-name, 7 | logging-format-interpolation, 8 | missing-docstring, 9 | no-self-use, 10 | not-callable, 11 | too-few-public-methods, 12 | too-many-arguments, 13 | too-many-instance-attributes, 14 | too-many-locals, 15 | unbalanced-tuple-unpacking, 16 | ungrouped-imports, 17 | unpacking-non-sequence, 18 | unused-argument, 19 | wrong-import-order, 20 | attribute-defined-outside-init, 21 | import-outside-toplevel, 22 | duplicate-code, 23 | logging-fstring-interpolation, 24 | import-error, 25 | no-name-in-module, 26 | dangerous-default-value, 27 | abstract-method, 28 | unspecified-encoding 29 | 30 | max-line-length = 120 31 | ignore-docstrings = yes 32 | ignored-modules = numpy,torch,cv2,openvino,onnx.onnx_cpp2py_export 33 | extension-pkg-whitelist = torch,cv2,scipy 34 | 35 | [SIMILARITIES] 36 | ignore-imports = yes 37 | 38 | [BASIC] 39 | bad-functions = print 40 | good-names = logger,fn 41 | 42 | [DESIGN] 43 | max-statements=80 44 | max-branches=20 45 | -------------------------------------------------------------------------------- /.style.yapf: -------------------------------------------------------------------------------- 1 | [style] 2 | BASED_ON_STYLE = pep8 3 | BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true 4 | SPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN = true 5 | DEDENT_CLOSING_BRACKETS = true 6 | SPACES_BEFORE_COMMENT = 1 7 | ARITHMETIC_PRECEDENCE_INDICATION = true -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | All notable changes to this project will be documented in this file. 4 | 5 | This project is based on [deep-person-reid project by KaiyangZhou](https://github.com/KaiyangZhou/deep-person-reid). 6 | With respect to it we made the following changes. 7 | 8 | ## \[2021-12-27\] 9 | ### Added 10 | * ImageAMSoftmaxEngine, MultilabelEngine with new features: mutual learning, am-softmax loss support, metric losses, EMA, SAM support. 11 | * Classification models, datasets, adapted engines for classification pipelines 12 | * Onnx and OpenVINO export 13 | * Multilabel classification support 14 | * Validation script 15 | * Multilabel classification metrics 16 | * Augmentations, optimizers, losses, learning rate schedulers 17 | * NNCF support 18 | 19 | ### Removed 20 | * Reid backbones and datasets 21 | 22 | ### Deprecated 23 | * Video data support 24 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (C) 2020-2021 Intel Corporation 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | 15 | 16 | MIT License 17 | 18 | Copyright (c) 2018 Kaiyang Zhou 19 | 20 | Permission is hereby granted, free of charge, to any person obtaining a copy 21 | of this software and associated documentation files (the "Software"), to deal 22 | in the Software without restriction, including without limitation the rights 23 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 24 | copies of the Software, and to permit persons to whom the Software is 25 | furnished to do so, subject to the following conditions: 26 | 27 | The above copyright notice and this permission notice shall be included in all 28 | copies or substantial portions of the Software. 29 | 30 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 31 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 32 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 33 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 34 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 35 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 36 | SOFTWARE. 37 | -------------------------------------------------------------------------------- /configs/classification/efficientnet_v1_b0.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: True 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 6 6 | step: 0.001 7 | smooth_f: 0.01 8 | epochs_warmup: 1 9 | path_to_savefig: 'lr_finder.jpg' 10 | max_lr: 0.01 11 | min_lr: 0.001 12 | n_trials: 15 13 | 14 | model: 15 | name: 'efficientnet_b0' 16 | type: 'classification' 17 | pretrained: True 18 | save_all_chkpts: False 19 | 20 | custom_datasets: 21 | roots: ['data/CIFAR100/train', 'data/CIFAR100/val'] 22 | types: ['classification_image_folder', 'classification_image_folder'] 23 | 24 | data: 25 | root: './' 26 | height: 224 27 | width: 224 28 | norm_mean: [0.485, 0.456, 0.406] 29 | norm_std: [0.229, 0.224, 0.225] 30 | save_dir: 'output/efficient_b0/log' 31 | workers: 6 32 | transforms: 33 | random_flip: 34 | enable: True 35 | p: 0.5 36 | random_rotate: 37 | enable: True 38 | p: 0.35 39 | angle: (-10,10) 40 | augmix: 41 | enable: True 42 | cfg_str: "augmix-m5-w3" 43 | cutout: 44 | enable: True 45 | cutout_factor: 0.35 46 | p: 0.35 47 | 48 | 49 | loss: 50 | name: 'am_softmax' 51 | softmax: 52 | compute_s: True 53 | 54 | sampler: 55 | train_sampler: 'RandomSampler' 56 | 57 | train: 58 | optim: 'sam' 59 | lr: 0.003 60 | nbd: True 61 | max_epoch: 200 62 | weight_decay: 5e-4 63 | batch_size: 64 64 | fixbase_epoch: 0 65 | lr_scheduler: 'warmup' 66 | warmup: 5 67 | base_scheduler: 'reduce_on_plateau' 68 | early_stopping: True 69 | train_patience: 5 70 | lr_decay_factor: 200 71 | deterministic: True 72 | patience: 5 73 | gamma: 0.1 74 | sam: 75 | rho: 0.05 76 | ema: 77 | enable: True 78 | ema_decay: 0.999 79 | mix_precision: True 80 | 81 | test: 82 | batch_size: 128 83 | evaluate: False 84 | eval_freq: 1 85 | -------------------------------------------------------------------------------- /configs/classification/efficientnet_v2_b0.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: True 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 6 6 | step: 0.001 7 | smooth_f: 0.01 8 | epochs_warmup: 1 9 | path_to_savefig: 'lr_finder.jpg' 10 | max_lr: 0.03 11 | min_lr: 0.003 12 | n_trials: 15 13 | 14 | model: 15 | name: 'efficientnetv2_b0' 16 | type: 'classification' 17 | pretrained: True 18 | save_all_chkpts: False 19 | 20 | custom_datasets: 21 | roots: ['data/CIFAR100/train', 'data/CIFAR100/val'] 22 | types: ['classification_image_folder', 'classification_image_folder'] 23 | 24 | data: 25 | root: './' 26 | height: 224 27 | width: 224 28 | norm_mean: [0.485, 0.456, 0.406] 29 | norm_std: [0.229, 0.224, 0.225] 30 | save_dir: 'output/efficientv2_b0/log' 31 | workers: 6 32 | transforms: 33 | random_flip: 34 | enable: True 35 | p: 0.5 36 | random_rotate: 37 | enable: True 38 | p: 0.35 39 | angle: (-10,10) 40 | augmix: 41 | enable: True 42 | cfg_str: "augmix-m5-w3" 43 | cutout: 44 | enable: True 45 | cutout_factor: 0.35 46 | p: 0.35 47 | 48 | loss: 49 | name: 'am_softmax' 50 | softmax: 51 | compute_s: True 52 | 53 | sampler: 54 | train_sampler: 'RandomSampler' 55 | 56 | train: 57 | optim: 'sam' 58 | lr: 0.003 59 | nbd: True 60 | max_epoch: 200 61 | weight_decay: 5e-4 62 | batch_size: 84 63 | fixbase_epoch: 0 64 | lr_scheduler: 'warmup' 65 | warmup: 5 66 | base_scheduler: 'reduce_on_plateau' 67 | early_stopping: True 68 | train_patience: 5 69 | lr_decay_factor: 200 70 | deterministic: True 71 | patience: 5 72 | gamma: 0.1 73 | sam: 74 | rho: 0.05 75 | ema: 76 | enable: True 77 | ema_decay: 0.999 78 | mix_precision: True 79 | 80 | test: 81 | batch_size: 128 82 | evaluate: False 83 | eval_freq: 1 84 | -------------------------------------------------------------------------------- /configs/classification/efficientnet_v2_small.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: True 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 6 6 | step: 0.001 7 | smooth_f: 0.01 8 | epochs_warmup: 1 9 | path_to_savefig: 'lr_finder.jpg' 10 | max_lr: 0.03 11 | min_lr: 0.003 12 | n_trials: 15 13 | 14 | model: 15 | name: 'efficientnetv2_s_21k' 16 | type: 'classification' 17 | pretrained: True 18 | save_all_chkpts: False 19 | 20 | custom_datasets: 21 | roots: ['data/CIFAR100/train', 'data/CIFAR100/val'] 22 | types: ['classification_image_folder', 'classification_image_folder'] 23 | 24 | data: 25 | root: './' 26 | height: 224 27 | width: 224 28 | norm_mean: [0.5, 0.5, 0.5] 29 | norm_std: [0.5, 0.5, 0.5] 30 | save_dir: 'output/efficientv2_s_21k/log' 31 | workers: 6 32 | transforms: 33 | random_flip: 34 | enable: True 35 | p: 0.5 36 | random_rotate: 37 | enable: True 38 | p: 0.35 39 | angle: (-10,10) 40 | augmix: 41 | enable: True 42 | cfg_str: "augmix-m5-w3" 43 | cutout: 44 | enable: True 45 | cutout_factor: 0.35 46 | p: 0.35 47 | 48 | loss: 49 | name: 'am_softmax' 50 | softmax: 51 | compute_s: True 52 | 53 | sampler: 54 | train_sampler: 'RandomSampler' 55 | 56 | 57 | train: 58 | optim: 'sam' 59 | lr: 0.005 60 | nbd: True 61 | max_epoch: 200 62 | weight_decay: 5e-4 63 | batch_size: 64 64 | fixbase_epoch: 0 65 | lr_scheduler: 'warmup' 66 | warmup: 5 67 | base_scheduler: 'reduce_on_plateau' 68 | early_stopping: True 69 | train_patience: 5 70 | lr_decay_factor: 200 71 | deterministic: True 72 | patience: 5 73 | gamma: 0.1 74 | sam: 75 | rho: 0.05 76 | ema: 77 | enable: True 78 | ema_decay: 0.999 79 | mix_precision: True 80 | 81 | test: 82 | batch_size: 128 83 | evaluate: False 84 | eval_freq: 1 85 | -------------------------------------------------------------------------------- /configs/classification/mobilenetv3_large.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: True 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 6 6 | step: 0.001 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 0.029 10 | min_lr: 0.005 11 | n_trials: 15 12 | 13 | model: 14 | name: 'mobilenetv3_large' 15 | type: 'classification' 16 | pretrained: True 17 | save_all_chkpts: True 18 | feature_dim: 1280 19 | 20 | mutual_learning: 21 | aux_configs: ['configs/classification_experiments/mobilenetv3_large_aux.yml'] 22 | 23 | custom_datasets: 24 | roots: ['data/CIFAR100/train', 'data/CIFAR100/val'] 25 | types: ['classification_image_folder', 'classification_image_folder'] 26 | 27 | data: 28 | root: './' 29 | height: 224 30 | width: 224 31 | norm_mean: [0.485, 0.456, 0.406] 32 | norm_std: [0.229, 0.224, 0.225] 33 | save_dir: 'output/mobilenetv3_large/log' 34 | workers: 6 35 | transforms: 36 | random_flip: 37 | enable: True 38 | p: 0.5 39 | random_rotate: 40 | enable: True 41 | p: 0.35 42 | angle: (-10,10) 43 | cutout: 44 | enable: True 45 | cutout_factor: 0.35 46 | p: 0.35 47 | augmix: 48 | enable: True 49 | cfg_str: "augmix-m5-w3" 50 | 51 | loss: 52 | name: 'softmax' 53 | softmax: 54 | s: 1.0 55 | compute_s: False 56 | 57 | sampler: 58 | train_sampler: 'RandomSampler' 59 | 60 | train: 61 | optim: 'sam' 62 | lr: 0.013 63 | nbd: True 64 | max_epoch: 200 65 | weight_decay: 5e-4 66 | batch_size: 84 67 | lr_scheduler: 'warmup' 68 | warmup: 5 69 | base_scheduler: 'reduce_on_plateau' 70 | early_stopping: True 71 | train_patience: 5 72 | lr_decay_factor: 200 73 | deterministic: True 74 | patience: 5 75 | gamma: 0.1 76 | sam: 77 | rho: 0.2 78 | ema: 79 | enable: True 80 | ema_decay: 0.999 81 | mix_precision: True 82 | 83 | test: 84 | batch_size: 128 85 | evaluate: False 86 | eval_freq: 1 87 | -------------------------------------------------------------------------------- /configs/classification/mobilenetv3_large_075.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: True 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 6 6 | step: 0.001 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 0.029 10 | min_lr: 0.005 11 | n_trials: 15 12 | 13 | model: 14 | name: 'mobilenetv3_large_075' 15 | type: 'classification' 16 | pretrained: True 17 | save_all_chkpts: True 18 | feature_dim: 1280 19 | 20 | mutual_learning: 21 | aux_configs: ['configs/classification/mobilenetv3_large_aux.yml'] 22 | 23 | custom_datasets: 24 | roots: ['data/CIFAR100/train', 'data/CIFAR100/val'] 25 | types: ['classification_image_folder', 'classification_image_folder'] 26 | 27 | data: 28 | root: './' 29 | height: 224 30 | width: 224 31 | norm_mean: [0.485, 0.456, 0.406] 32 | norm_std: [0.229, 0.224, 0.225] 33 | save_dir: 'output/mobilenetv3_large_075/log' 34 | workers: 6 35 | transforms: 36 | random_flip: 37 | enable: True 38 | p: 0.5 39 | random_rotate: 40 | enable: True 41 | p: 0.35 42 | angle: (-10,10) 43 | cutout: 44 | enable: True 45 | cutout_factor: 0.35 46 | p: 0.35 47 | augmix: 48 | enable: True 49 | cfg_str: "augmix-m4-w2" 50 | 51 | loss: 52 | name: 'softmax' 53 | softmax: 54 | s: 1.0 55 | compute_s: False 56 | 57 | sampler: 58 | train_sampler: 'RandomSampler' 59 | 60 | train: 61 | optim: 'sam' 62 | lr: 0.013 63 | nbd: True 64 | max_epoch: 200 65 | weight_decay: 5e-4 66 | batch_size: 100 67 | lr_scheduler: 'warmup' 68 | warmup: 5 69 | base_scheduler: 'reduce_on_plateau' 70 | early_stopping: True 71 | train_patience: 5 72 | lr_decay_factor: 200 73 | deterministic: True 74 | patience: 5 75 | gamma: 0.1 76 | sam: 77 | rho: 0.15 78 | ema: 79 | enable: True 80 | ema_decay: 0.999 81 | mix_precision: True 82 | 83 | test: 84 | batch_size: 128 85 | evaluate: False 86 | eval_freq: 1 87 | -------------------------------------------------------------------------------- /configs/classification/mobilenetv3_large_aux.yml: -------------------------------------------------------------------------------- 1 | model: 2 | name: 'mobilenetv3_large' 3 | type: 'classification' 4 | pretrained: True 5 | feature_dim: 1280 6 | dropout_cls: 7 | p: 0.2 8 | dist: 'bernoulli' 9 | 10 | loss: 11 | name: 'am_softmax' 12 | softmax: 13 | s: 1.0 14 | compute_s: False 15 | 16 | 17 | sampler: 18 | train_sampler: 'RandomSampler' 19 | 20 | train: 21 | optim: 'sam' 22 | lr: 0.013 23 | nbd: True 24 | weight_decay: 5e-4 25 | lr_scheduler: 'warmup' 26 | warmup: 5 27 | base_scheduler: 'reduce_on_plateau' 28 | early_stopping: True 29 | train_patience: 5 30 | lr_decay_factor: 200 31 | deterministic: True 32 | patience: 5 33 | gamma: 0.1 34 | sam: 35 | rho: 0.05 36 | mix_precision: True 37 | 38 | test: 39 | batch_size: 128 40 | evaluate: False 41 | eval_freq: 1 42 | -------------------------------------------------------------------------------- /configs/classification/mobilenetv3_small.yml: -------------------------------------------------------------------------------- 1 | 2 | lr_finder: 3 | enable: True 4 | mode: TPE 5 | stop_after: False 6 | num_epochs: 6 7 | step: 0.001 8 | epochs_warmup: 1 9 | path_to_savefig: 'lr_finder.jpg' 10 | max_lr: 0.029 11 | min_lr: 0.005 12 | n_trials: 15 13 | 14 | model: 15 | name: 'mobilenetv3_small' 16 | type: 'classification' 17 | pretrained: True 18 | save_all_chkpts: False 19 | 20 | mutual_learning: 21 | aux_configs: ['configs/classification/mobilenetv3_large_aux.yml'] 22 | 23 | custom_datasets: 24 | roots: ['data/CIFAR100/train', 'data/CIFAR100/val'] 25 | types: ['classification_image_folder', 'classification_image_folder'] 26 | 27 | data: 28 | root: './' 29 | height: 224 30 | width: 224 31 | norm_mean: [0.485, 0.456, 0.406] 32 | norm_std: [0.229, 0.224, 0.225] 33 | save_dir: 'output/mobilenetv3_small/log' 34 | workers: 6 35 | transforms: 36 | random_flip: 37 | enable: True 38 | p: 0.5 39 | random_rotate: 40 | enable: True 41 | p: 0.35 42 | angle: (-10,10) 43 | augmix: 44 | enable: True 45 | cfg_str: "augmix-m4-w2" 46 | cutout: 47 | enable: True 48 | cutout_factor: 0.35 49 | p: 0.35 50 | 51 | loss: 52 | name: 'softmax' 53 | softmax: 54 | s: 1.0 55 | compute_s: False 56 | 57 | sampler: 58 | train_sampler: 'RandomSampler' 59 | 60 | train: 61 | optim: 'sam' 62 | lr: 0.013 63 | nbd: True 64 | max_epoch: 200 65 | weight_decay: 5e-4 66 | batch_size: 128 67 | lr_scheduler: 'warmup' 68 | warmup: 5 69 | base_scheduler: 'reduce_on_plateau' 70 | early_stopping: True 71 | train_patience: 5 72 | lr_decay_factor: 200 73 | deterministic: True 74 | patience: 5 75 | gamma: 0.1 76 | sam: 77 | rho: 0.07 78 | ema: 79 | enable: True 80 | ema_decay: 0.999 81 | mix_precision: True 82 | 83 | test: 84 | batch_size: 128 85 | evaluate: False 86 | eval_freq: 1 87 | -------------------------------------------------------------------------------- /configs/multilabel_classification/efficientnet_v1_b0.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: False 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 10 6 | step: 1e-5 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 1e-3 10 | min_lr: 1e-5 11 | n_trials: 15 12 | 13 | model: 14 | name: 'efficientnet_b0' 15 | type: 'multilabel' 16 | pretrained: True 17 | save_all_chkpts: False 18 | dropout_cls: 19 | p: 0.1 20 | 21 | custom_datasets: 22 | roots: ['data/coco/train.json', 'data/coco/val.json'] 23 | types: ['multilabel_classification', 'multilabel_classification'] 24 | 25 | data: 26 | root: './' 27 | height: 448 28 | width: 448 29 | norm_mean: [0.485, 0.456, 0.406] 30 | norm_std: [0.229, 0.224, 0.225] 31 | save_dir: 'output/mulitilabel/efficientnet_b0' 32 | workers: 6 33 | transforms: 34 | random_flip: 35 | enable: True 36 | p: 0.5 37 | randaugment: 38 | enable: True 39 | cutout: 40 | enable: True 41 | cutout_factor: 0.35 42 | p: 0.35 43 | 44 | loss: 45 | name: 'asl' 46 | softmax: 47 | s: 1.0 48 | compute_s: False 49 | asl: 50 | gamma_pos: 0. 51 | gamma_neg: 4. 52 | 53 | sampler: 54 | train_sampler: 'RandomSampler' 55 | 56 | train: 57 | optim: 'sam' 58 | lr: 0.0002 59 | nbd: True 60 | max_epoch: 80 61 | weight_decay: 5e-4 62 | batch_size: 64 63 | lr_scheduler: 'onecycle' 64 | early_stopping: True 65 | pct_start: 0.2 66 | train_patience: 5 67 | lr_decay_factor: 1000 68 | deterministic: True 69 | target_metric: test_acc 70 | gamma: 0.1 71 | sam: 72 | rho: 0.05 73 | ema: 74 | enable: True 75 | ema_decay: 0.9995 76 | mix_precision: True 77 | 78 | test: 79 | batch_size: 128 80 | evaluate: False 81 | eval_freq: 1 82 | -------------------------------------------------------------------------------- /configs/multilabel_classification/efficientnet_v2_b0.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: False 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 10 6 | step: 1e-5 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 1e-3 10 | min_lr: 1e-5 11 | n_trials: 15 12 | 13 | model: 14 | name: 'efficientnetv2_b0' 15 | type: 'multilabel' 16 | pretrained: True 17 | save_all_chkpts: False 18 | dropout_cls: 19 | p: 0.1 20 | 21 | custom_datasets: 22 | roots: ['data/coco/train.json', 'data/coco/val.json'] 23 | types: ['multilabel_classification', 'multilabel_classification'] 24 | 25 | data: 26 | root: './' 27 | height: 448 28 | width: 448 29 | norm_mean: [0.485, 0.456, 0.406] 30 | norm_std: [0.229, 0.224, 0.225] 31 | save_dir: 'output/mulitilabel/efficientnetv2_b0' 32 | workers: 6 33 | transforms: 34 | random_flip: 35 | enable: True 36 | p: 0.5 37 | randaugment: 38 | enable: True 39 | cutout: 40 | enable: True 41 | cutout_factor: 0.35 42 | p: 0.35 43 | 44 | loss: 45 | name: 'asl' 46 | softmax: 47 | s: 1.0 48 | compute_s: False 49 | asl: 50 | gamma_pos: 0. 51 | gamma_neg: 4. 52 | 53 | sampler: 54 | train_sampler: 'RandomSampler' 55 | 56 | train: 57 | optim: 'sam' 58 | lr: 0.0002 59 | nbd: True 60 | max_epoch: 80 61 | weight_decay: 5e-4 62 | batch_size: 64 63 | lr_scheduler: 'onecycle' 64 | early_stopping: True 65 | pct_start: 0.2 66 | train_patience: 5 67 | lr_decay_factor: 1000 68 | deterministic: True 69 | target_metric: test_acc 70 | gamma: 0.1 71 | sam: 72 | rho: 0.05 73 | ema: 74 | enable: True 75 | ema_decay: 0.9995 76 | mix_precision: True 77 | 78 | test: 79 | batch_size: 128 80 | evaluate: False 81 | eval_freq: 1 82 | -------------------------------------------------------------------------------- /configs/multilabel_classification/efficientnet_v2_small.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: False 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 10 6 | step: 1e-5 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 1e-3 10 | min_lr: 1e-5 11 | n_trials: 15 12 | 13 | model: 14 | name: 'efficientnetv2_s_21k' 15 | type: 'multilabel' 16 | pretrained: True 17 | save_all_chkpts: False 18 | dropout_cls: 19 | p: 0.1 20 | 21 | custom_datasets: 22 | roots: ['data/coco/train.json', 'data/coco/val.json'] 23 | types: ['multilabel_classification', 'multilabel_classification'] 24 | 25 | data: 26 | root: './' 27 | height: 448 28 | width: 448 29 | norm_mean: [0.5, 0.5, 0.5] 30 | norm_std: [0.5, 0.5, 0.5] 31 | save_dir: 'output/mulitilabel/efficientnetv2_s_21k' 32 | workers: 6 33 | transforms: 34 | random_flip: 35 | enable: True 36 | p: 0.5 37 | randaugment: 38 | enable: True 39 | cutout: 40 | enable: True 41 | cutout_factor: 0.35 42 | p: 0.35 43 | 44 | loss: 45 | name: 'asl' 46 | softmax: 47 | s: 1.0 48 | compute_s: False 49 | asl: 50 | gamma_pos: 0. 51 | gamma_neg: 4. 52 | 53 | sampler: 54 | train_sampler: 'RandomSampler' 55 | 56 | train: 57 | optim: 'sam' 58 | lr: 0.0002 59 | nbd: True 60 | max_epoch: 80 61 | weight_decay: 5e-4 62 | batch_size: 64 63 | lr_scheduler: 'onecycle' 64 | early_stopping: True 65 | pct_start: 0.2 66 | train_patience: 5 67 | lr_decay_factor: 1000 68 | deterministic: True 69 | target_metric: test_acc 70 | gamma: 0.1 71 | sam: 72 | rho: 0.05 73 | ema: 74 | enable: True 75 | ema_decay: 0.9995 76 | mix_precision: True 77 | 78 | test: 79 | batch_size: 128 80 | evaluate: False 81 | eval_freq: 1 82 | -------------------------------------------------------------------------------- /configs/multilabel_classification/mobilenetv3_large.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: False 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 7 6 | step: 0.001 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 1e-3 10 | min_lr: 1e-5 11 | n_trials: 15 12 | 13 | model: 14 | name: 'mobilenetv3_large_21k' 15 | type: 'multilabel' 16 | pretrained: True 17 | save_all_chkpts: False 18 | dropout_cls: 19 | p: 0.1 20 | 21 | custom_datasets: 22 | roots: ['data/coco/train.json', 'data/coco/val.json'] 23 | types: ['multilabel_classification', 'multilabel_classification'] 24 | 25 | data: 26 | root: './' 27 | height: 448 28 | width: 448 29 | norm_mean: [0, 0, 0] 30 | norm_std: [1, 1, 1] 31 | save_dir: 'output/mulitilabel/mobilenetv3_large_21k' 32 | workers: 6 33 | transforms: 34 | random_flip: 35 | enable: True 36 | p: 0.5 37 | randaugment: 38 | enable: True 39 | cutout: 40 | enable: True 41 | cutout_factor: 0.35 42 | p: 0.35 43 | 44 | loss: 45 | name: 'asl' 46 | softmax: 47 | s: 1.0 48 | compute_s: False 49 | asl: 50 | gamma_pos: 0. 51 | gamma_neg: 4. 52 | 53 | sampler: 54 | train_sampler: 'RandomSampler' 55 | 56 | train: 57 | optim: 'sam' 58 | lr: 0.0002 59 | nbd: True 60 | max_epoch: 80 61 | weight_decay: 5e-4 62 | batch_size: 64 63 | lr_scheduler: 'onecycle' 64 | early_stopping: True 65 | pct_start: 0.2 66 | train_patience: 5 67 | lr_decay_factor: 1000 68 | deterministic: True 69 | target_metric: test_acc 70 | gamma: 0.1 71 | sam: 72 | rho: 0.05 73 | ema: 74 | enable: True 75 | ema_decay: 0.9997 76 | mix_precision: True 77 | 78 | test: 79 | batch_size: 128 80 | evaluate: False 81 | eval_freq: 1 82 | -------------------------------------------------------------------------------- /configs/multilabel_classification/mobilenetv3_large_75.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: True 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 7 6 | step: 0.001 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 1e-3 10 | min_lr: 1e-5 11 | n_trials: 15 12 | 13 | model: 14 | name: 'mobilenetv3_large_075' 15 | type: 'multilabel' 16 | pretrained: True 17 | save_all_chkpts: False 18 | dropout_cls: 19 | p: 0.1 20 | 21 | custom_datasets: 22 | roots: ['data/coco/train.json', 'data/coco/val.json'] 23 | types: ['multilabel_classification', 'multilabel_classification'] 24 | 25 | data: 26 | root: './' 27 | height: 448 28 | width: 448 29 | norm_mean: [0.485, 0.456, 0.406] 30 | norm_std: [0.229, 0.224, 0.225] 31 | save_dir: 'output/mulitilabel/mobilenetv3_large_75' 32 | workers: 6 33 | transforms: 34 | random_flip: 35 | enable: True 36 | p: 0.5 37 | randaugment: 38 | enable: True 39 | cutout: 40 | enable: True 41 | cutout_factor: 0.35 42 | p: 0.35 43 | 44 | loss: 45 | name: 'asl' 46 | softmax: 47 | s: 1.0 48 | compute_s: False 49 | asl: 50 | gamma_pos: 0. 51 | gamma_neg: 4. 52 | 53 | sampler: 54 | train_sampler: 'RandomSampler' 55 | 56 | 57 | train: 58 | optim: 'sam' 59 | lr: 0.0002 60 | nbd: True 61 | max_epoch: 80 62 | weight_decay: 5e-4 63 | batch_size: 64 64 | lr_scheduler: 'onecycle' 65 | early_stopping: True 66 | pct_start: 0.2 67 | train_patience: 5 68 | lr_decay_factor: 1000 69 | deterministic: True 70 | target_metric: test_acc 71 | gamma: 0.1 72 | sam: 73 | rho: 0.05 74 | ema: 75 | enable: True 76 | ema_decay: 0.9995 77 | mix_precision: True 78 | 79 | test: 80 | batch_size: 128 81 | evaluate: False 82 | eval_freq: 1 83 | -------------------------------------------------------------------------------- /configs/multilabel_classification/mobilenetv3_small.yml: -------------------------------------------------------------------------------- 1 | lr_finder: 2 | enable: False 3 | mode: TPE 4 | stop_after: False 5 | num_epochs: 7 6 | step: 0.001 7 | epochs_warmup: 1 8 | path_to_savefig: 'lr_finder.jpg' 9 | max_lr: 1e-3 10 | min_lr: 1e-5 11 | n_trials: 15 12 | 13 | model: 14 | name: 'mobilenetv3_small' 15 | type: 'multilabel' 16 | pretrained: True 17 | save_all_chkpts: False 18 | 19 | custom_datasets: 20 | roots: ['data/coco/train.json', 'data/coco/val.json'] 21 | types: ['multilabel_classification', 'multilabel_classification'] 22 | 23 | data: 24 | root: './' 25 | height: 448 26 | width: 448 27 | norm_mean: [0.485, 0.456, 0.406] 28 | norm_std: [0.229, 0.224, 0.225] 29 | save_dir: 'output/mulitilabel/mobilenetv3_small' 30 | workers: 6 31 | transforms: 32 | random_flip: 33 | enable: True 34 | p: 0.5 35 | randaugment: 36 | enable: True 37 | cutout: 38 | enable: True 39 | cutout_factor: 0.35 40 | p: 0.35 41 | 42 | loss: 43 | name: 'asl' 44 | softmax: 45 | s: 1.0 46 | compute_s: False 47 | asl: 48 | gamma_pos: 0. 49 | gamma_neg: 4. 50 | 51 | sampler: 52 | train_sampler: 'RandomSampler' 53 | 54 | train: 55 | optim: 'sam' 56 | lr: 0.0002 57 | nbd: True 58 | max_epoch: 80 59 | weight_decay: 5e-4 60 | batch_size: 64 61 | lr_scheduler: 'onecycle' 62 | early_stopping: True 63 | pct_start: 0.2 64 | train_patience: 5 65 | lr_decay_factor: 1000 66 | deterministic: True 67 | target_metric: test_acc 68 | gamma: 0.1 69 | sam: 70 | rho: 0.05 71 | ema: 72 | enable: True 73 | ema_decay: 0.9995 74 | mix_precision: True 75 | 76 | test: 77 | batch_size: 128 78 | evaluate: False 79 | eval_freq: 1 80 | -------------------------------------------------------------------------------- /configs/opt_configs/example.json: -------------------------------------------------------------------------------- 1 | 2 | { 3 | "epochs" : 1, 4 | "n_trials": 1, 5 | "float": 6 | [ 7 | { 8 | "name": "train.lr", 9 | "range": [0.001, 0.1], 10 | "step": 0.001 11 | }, 12 | { 13 | "name": "data.transforms.cutout.cutout_factor", 14 | "range": [0.1, 1.0], 15 | "step": 0.1 16 | } 17 | ], 18 | "int":[ 19 | { 20 | "name": "loss.softmax.s", 21 | "range": [1.0, 50.0], 22 | "step": 1.0 23 | } 24 | ], 25 | "categorical":[ 26 | { 27 | "name": "model.pooling_type", 28 | "range": ["avg","max","avg+max"] 29 | } 30 | ] 31 | 32 | } 33 | -------------------------------------------------------------------------------- /constraints.txt: -------------------------------------------------------------------------------- 1 | pytest==6.2.5 2 | networkx==2.6.3 3 | -------------------------------------------------------------------------------- /init_venv.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -v 3 | set -x 4 | 5 | work_dir=$(realpath "$(dirname $0)") 6 | 7 | venv_dir=$1 8 | PYTHON_NAME=$2 9 | 10 | if [ -z "$venv_dir" ]; then 11 | venv_dir=$(realpath -m ${work_dir}/venv) 12 | else 13 | venv_dir=$(realpath -m "$venv_dir") 14 | fi 15 | 16 | if [[ -z $PYTHON_NAME ]]; then 17 | # the default option -- note that the minimal version of 18 | # python that is suitable for this repo is python3.7, 19 | # whereas the default python3 may point to python3.6 20 | PYTHON_NAME=python3 21 | fi 22 | 23 | PYTHON_VERSION=$($PYTHON_NAME --version | sed -e "s/^Python \([0-9]\.[0-9]\)\..*/\1/") || exit 1 24 | if [[ $PYTHON_VERSION != "3.7" && $PYTHON_VERSION != "3.8" && $PYTHON_VERSION != "3.9" ]]; then 25 | echo "Wrong version of python: '$PYTHON_VERSION'" 26 | exit 1 27 | fi 28 | 29 | cd ${work_dir} 30 | 31 | if [[ -e ${venv_dir} ]]; then 32 | echo 33 | echo "Virtualenv already exists. Use command to start working:" 34 | echo "$ . ${venv_dir}/bin/activate" 35 | exit 36 | fi 37 | 38 | # Create virtual environment 39 | $PYTHON_NAME -m venv ${venv_dir} --prompt="classification" 40 | 41 | if ! [ -e "${venv_dir}/bin/activate" ]; then 42 | echo "The virtual environment was not created." 43 | exit 44 | fi 45 | 46 | . ${venv_dir}/bin/activate 47 | 48 | # Get CUDA version. 49 | CUDA_HOME_CANDIDATE=/usr/local/cuda 50 | if [ -z "${CUDA_HOME}" ] && [ -d ${CUDA_HOME_CANDIDATE} ]; then 51 | echo "Exporting CUDA_HOME as ${CUDA_HOME_CANDIDATE}" 52 | export CUDA_HOME=${CUDA_HOME_CANDIDATE} 53 | fi 54 | 55 | if [ -e "$CUDA_HOME" ]; then 56 | if [ -e "$CUDA_HOME/version.txt" ]; then 57 | # Get CUDA version from version.txt file. 58 | CUDA_VERSION=$(cat $CUDA_HOME/version.txt | sed -e "s/^.*CUDA Version *//" -e "s/ .*//") 59 | else 60 | # Get CUDA version from directory name. 61 | CUDA_HOME_DIR=`readlink -f $CUDA_HOME` 62 | CUDA_HOME_DIR=`basename $CUDA_HOME_DIR` 63 | CUDA_VERSION=`echo $CUDA_HOME_DIR | cut -d "-" -f 2` 64 | fi 65 | fi 66 | 67 | # install PyTorch and MMCV. 68 | export TORCH_VERSION=1.8.2 69 | export TORCHVISION_VERSION=0.9.2 70 | 71 | if [[ -z ${CUDA_VERSION} ]]; then 72 | echo "CUDA was not found, installing dependencies in CPU-only mode. If you want to use CUDA, set CUDA_HOME and CUDA_VERSION beforehand." 73 | else 74 | # Remove dots from CUDA version string, if any. 75 | CUDA_VERSION_CODE=$(echo ${CUDA_VERSION} | sed -e "s/\.//" -e "s/\(...\).*/\1/") 76 | echo "Using CUDA_VERSION ${CUDA_VERSION}" 77 | if [[ "${CUDA_VERSION_CODE}" != "111" ]] && [[ "${CUDA_VERSION_CODE}" != "102" ]] ; then 78 | echo "CUDA version must be either 11.1 or 10.2" 79 | exit 1 80 | fi 81 | echo "export CUDA_HOME=${CUDA_HOME}" >> ${venv_dir}/bin/activate 82 | fi 83 | 84 | CONSTRAINTS_FILE=$(tempfile) 85 | export PIP_CONSTRAINT=${CONSTRAINTS_FILE} 86 | 87 | # Newer versions of pip have troubles with NNCF installation from the repo commit. 88 | pip install pip==21.2.1 || exit 1 89 | 90 | if [[ -z $CUDA_VERSION_CODE ]]; then 91 | export TORCH_VERSION=${TORCH_VERSION}+cpu 92 | export TORCHVISION_VERSION=${TORCHVISION_VERSION}+cpu 93 | else 94 | export TORCH_VERSION=${TORCH_VERSION}+cu${CUDA_VERSION_CODE} 95 | export TORCHVISION_VERSION=${TORCHVISION_VERSION}+cu${CUDA_VERSION_CODE} 96 | fi 97 | 98 | pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html --no-cache || exit 1 99 | echo torch==${TORCH_VERSION} >> ${CONSTRAINTS_FILE} 100 | echo torchvision==${TORCHVISION_VERSION} >> ${CONSTRAINTS_FILE} 101 | 102 | # Install other requirements. 103 | pip install -r requirements.txt -c ${CONSTRAINTS_FILE} || exit 1 104 | cat openvino-requirements.txt | xargs -n 1 -L 1 pip install -c ${CONSTRAINTS_FILE} || exit 1 105 | 106 | pip install -e . || exit 1 107 | 108 | if [[ ! -z $OTE_SDK_PATH ]]; then 109 | pip install -e $OTE_SDK_PATH || exit 1 110 | elif [[ ! -z $SC_SDK_REPO ]]; then 111 | pip install -e $SC_SDK_REPO/src/ote_sdk || exit 1 112 | else 113 | echo "OTE_SDK_PATH or SC_SDK_REPO should be specified" 114 | exit 1 115 | fi 116 | 117 | 118 | deactivate 119 | 120 | echo 121 | echo "Activate a virtual environment to start working:" 122 | echo "$ . ${venv_dir}/bin/activate" 123 | -------------------------------------------------------------------------------- /linter.sh: -------------------------------------------------------------------------------- 1 | echo "Running isort" 2 | isort -y -sp . 3 | echo "Done" 4 | 5 | echo "Running yapf" 6 | yapf -i -r -vv -e build . 7 | echo "Done" 8 | 9 | echo "Running flake8" 10 | flake8 . 11 | echo "Done" -------------------------------------------------------------------------------- /openvino-requirements.txt: -------------------------------------------------------------------------------- 1 | openvino==2022.3.0 2 | openvino-dev==2022.3.0 3 | openmodelzoo-modelapi==2022.3.0 4 | nncf@ git+https://github.com/openvinotoolkit/nncf@e2848e58ca3608c9a7145d0b5f2b91277095b87a#egg=nncf 5 | -------------------------------------------------------------------------------- /optional-requirements.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | pylint==2.11.1 3 | flake8 4 | yapf 5 | isort -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | Pillow 3 | six 4 | scipy==1.8.1 5 | opencv-python 6 | matplotlib 7 | tb-nightly 8 | future 9 | yacs 10 | gdown 11 | scikit-learn 12 | terminaltables 13 | pytorchcv 14 | torch-lr-finder==0.2.1 15 | onnx 16 | torchvision==0.9.* 17 | torch>=1.13.1 18 | optuna==3.0.2 19 | timm>=0.6.5 20 | addict 21 | randaugment 22 | ptflops 23 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | import os.path as osp 9 | from setuptools import setup, Extension, find_packages 10 | 11 | import numpy as np 12 | 13 | repo_root = osp.dirname(osp.realpath(__file__)) 14 | 15 | def readme(): 16 | with open(osp.join(repo_root, 'README.rst')) as f: 17 | content = f.read() 18 | return content 19 | 20 | 21 | def find_version(): 22 | version_file = osp.join(repo_root, 'torchreid/version.py') 23 | with open(version_file, 'r') as f: 24 | exec(compile(f.read(), version_file, 'exec')) 25 | return locals()['__version__'] 26 | 27 | 28 | def numpy_include(): 29 | try: 30 | numpy_include = np.get_include() 31 | except AttributeError: 32 | numpy_include = np.get_numpy_include() 33 | return numpy_include 34 | 35 | 36 | def get_requirements(filename): 37 | requires = [] 38 | links = [] 39 | with open(osp.join(repo_root, filename), 'r') as f: 40 | for line in f.readlines(): 41 | line = line.replace('\n', '') 42 | if '-f http' in line: 43 | links.append(line) 44 | else: 45 | requires.append(line) 46 | return requires, links 47 | 48 | packages, links = get_requirements('requirements.txt') 49 | 50 | setup( 51 | name='torchreid', 52 | version=find_version(), 53 | description='A library for deep learning object re-ID and classification in PyTorch', 54 | author='Kaiyang Zhou, Intel Corporation', 55 | license='Apache-2.0', 56 | long_description=readme(), 57 | url='https://github.com/openvinotoolkit/deep-object-reid', 58 | dependency_links=links, 59 | packages=find_packages(), 60 | install_requires=packages, 61 | keywords=['Object Re-Identification', 'Image Classification', 'Deep Learning', 'Computer Vision'], 62 | ) 63 | -------------------------------------------------------------------------------- /third-party-programs.txt: -------------------------------------------------------------------------------- 1 | Deep Object Reid Third Party Programs File 2 | 3 | This file contains the list of third party software ("third party programs") 4 | contained in the Intel software and their required notices and/or license 5 | terms. This third party software, even if included with the distribution of 6 | the Intel software, may be governed by separate license terms, including 7 | without limitation, third party license terms, other Intel software license 8 | terms, and open source software license terms. These separate license terms 9 | govern your use of the third party programs as set forth in the 10 | "third-party-programs.txt" or other similarly-named text file. 11 | 12 | Third party programs and their corresponding required notices and/or license 13 | terms are listed below. 14 | 15 | ------------------------------------------------------------- 16 | 17 | 1. fmix 18 | Copyright (c) 2020 Vision, Learning and Control Research Group 19 | 20 | MIT License 21 | 22 | Copyright (c) 2020 Vision, Learning and Control Research Group 23 | 24 | Permission is hereby granted, free of charge, to any person obtaining a copy 25 | of this software and associated documentation files (the "Software"), to deal 26 | in the Software without restriction, including without limitation the rights 27 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 28 | copies of the Software, and to permit persons to whom the Software is 29 | furnished to do so, subject to the following conditions: 30 | 31 | The above copyright notice and this permission notice shall be included in all 32 | copies or substantial portions of the Software. 33 | 34 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 35 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 36 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 37 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 38 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 39 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 40 | SOFTWARE. 41 | 42 | ------------------------------------------------------------- 43 | 44 | 2. GeneralizedMeanPooling 45 | Copyright (c) 2019 mangye16 46 | 47 | ------------------------------------------------------------- 48 | 49 | 3. CenterCrop, CoarseDropout, Cutout, Equalize, GaussianBlur, Posterize, Posterize, RandomCrop 50 | Copyright (c) 2017 Buslaev Alexander, Alexander Parinov, Vladimir Iglovikov 51 | SPDX-License-Identifier: MIT 52 | 53 | ------------------------------------------------------------- 54 | 55 | 4. augment_and_mix_transform 56 | Copyright (c) 2019 Ross Wightman 57 | SPDX-License-Identifier: Apache-2.0 58 | 59 | ------------------------------------------------------------- 60 | 61 | 5. depthwise_conv3x3, get_activation_layer, Concurrent, ConvBlock, conv1x1_block, conv1x1, 62 | conv3x3_block, dwconv_block, dwconv3x3_block, dwconv5x5_block, DwsConvBlock, 63 | dwsconv3x3_block, round_channels, Identity, Swish, HSigmoid, HSwish, SEBlock 64 | Copyright (c) 2018-2021 Oleg Sémery 65 | SPDX-License-Identifier: MIT 66 | 67 | ------------------------------------------------------------- 68 | 69 | 6. File torchreid/models/efficient_net_pytcv.py 70 | Copyright (c) 2018-2021 Oleg Sémery 71 | SPDX-License-Identifier: MIT 72 | 73 | ------------------------------------------------------------- 74 | 75 | 7. File torchreid/models/mobilenetv3.py 76 | Copyright (c) 2019 Duo LI 77 | SPDX-License-Identifier: MIT 78 | 79 | ------------------------------------------------------------- 80 | 81 | 8. File torchreid/models/model_store.py 82 | Copyright (c) 2018-2021 Oleg Sémery 83 | SPDX-License-Identifier: MIT 84 | 85 | ------------------------------------------------------------- 86 | 87 | 9. StateCacher 88 | Copyright (c) 2018 davidtvs 89 | SPDX-License-Identifier: MIT 90 | 91 | ------------------------------------------------------------- 92 | 93 | 10. ModelEmaV2 94 | Copyright (c) 2019 Ross Wightman 95 | SPDX-License-Identifier: Apache-2.0 96 | 97 | ------------------------------------------------------------- 98 | 99 | 11. TripletLoss, Logger 100 | Copyright (c) 2017 Tong Xiao 101 | SPDX-License-Identifier: MIT 102 | 103 | ------------------------------------------------------------- 104 | 105 | 12. RandomErasing 106 | Copyright (c) 2018-2021 Zhun Zhong 107 | SPDX-License-Identifier: Apache-2.0 108 | 109 | ------------------------------------------------------------- 110 | 111 | 13. AugMixAugment 112 | Copyright 2019 Google LLC 113 | SPDX-License-Identifier: Apache-2.0 114 | 115 | ------------------------------------------------------------- 116 | 117 | 14. collect_env_info 118 | Copyright (c) 2018 Facebook 119 | SPDX-License-Identifier: MIT 120 | 121 | ------------------------------------------------------------- 122 | 123 | 15. RAdam 124 | Copyright (c) 2019 Liyuan Liu 125 | SPDX-License-Identifier: Apache-2.0 126 | 127 | ------------------------------------------------------------- 128 | 129 | 16. SAM 130 | Copyright 2020 Google Research 131 | SPDX-License-Identifier: Apache-2.0 132 | 133 | ------------------------------------------------------------- 134 | 135 | 17. make_divisible 136 | Copyright 2022 Google LLC. All rights reserved. 137 | SPDX-License-Identifier: Apache-2.0 138 | -------------------------------------------------------------------------------- /tools/auxiliary/add_initial_lr_to_checkpoint.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2020-2021 Intel Corporation 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | 5 | import argparse 6 | import os 7 | 8 | import torch 9 | 10 | from scripts.default_config import get_default_config 11 | from torchreid.utils import save_checkpoint 12 | 13 | def main(): 14 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, 15 | description='The script adds the initial LR value ' 16 | 'to deep-object-reid checkpoints ' 17 | '-- it will allow using it for NNCF, ' 18 | 'NNCF part will initialize its LR from the checkpoints`s LR') 19 | parser.add_argument('--lr', type=float, required=True, 20 | help='path to config file') 21 | parser.add_argument('--checkpoint', type=str, required=True, 22 | help='path to the src checkpoint file') 23 | parser.add_argument('--dst-folder', type=str, required=True, 24 | help='path to the dst folder to store dst checkpoint file') 25 | args = parser.parse_args() 26 | 27 | checkpoint = torch.load(args.checkpoint, map_location='cpu') 28 | if not isinstance(checkpoint, dict): 29 | raise RuntimeError('Wrong format of checkpoint -- it is not the result of deep-object-reid training') 30 | if checkpoint.get('initial_lr'): 31 | raise RuntimeError(f'Checkpoint {args.checkpoint} already has initial_lr') 32 | 33 | if not os.path.isdir(args.dst_folder): 34 | raise RuntimeError(f'The dst folder {args.dst_folder} is NOT present') 35 | 36 | checkpoint['initial_lr'] = float(args.lr) 37 | res_path = save_checkpoint(checkpoint, args.dst_folder) 38 | 39 | 40 | if __name__ == '__main__': 41 | main() 42 | -------------------------------------------------------------------------------- /tools/auxiliary/add_nncf_metainfo_to_checkpoint.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2020-2021 Intel Corporation 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | 5 | import argparse 6 | import os 7 | 8 | import torch 9 | 10 | from scripts.default_config import get_default_config, merge_from_files_with_base 11 | from torchreid.utils import save_checkpoint 12 | from torchreid.integration.nncf.compression import get_default_nncf_compression_config 13 | 14 | def main(): 15 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, 16 | description='The script adds the default int8 quantization NNCF metainfo ' 17 | 'to NNCF deep-object-reid checkpoints ' 18 | 'that were trained when NNCF metainfo was not ' 19 | 'stored in NNCF checkpoints') 20 | parser.add_argument('--config-file', type=str, required=True, 21 | help='path to config file') 22 | parser.add_argument('--checkpoint', type=str, required=True, 23 | help='path to the src checkpoint file') 24 | parser.add_argument('--dst-folder', type=str, required=True, 25 | help='path to the dst folder to store dst checkpoint file') 26 | args = parser.parse_args() 27 | 28 | cfg = get_default_config() 29 | merge_from_files_with_base(cfg, args.config_file) 30 | checkpoint = torch.load(args.checkpoint, map_location='cpu') 31 | if not isinstance(checkpoint, dict): 32 | raise RuntimeError('Wrong format of checkpoint -- it is not the result of deep-object-reid training') 33 | if checkpoint.get('nncf_metainfo'): 34 | raise RuntimeError(f'Checkpoint {args.checkpoint} already has nncf_metainfo') 35 | 36 | if not os.path.isdir(args.dst_folder): 37 | raise RuntimeError(f'The dst folder {args.dst_folder} is NOT present') 38 | 39 | # default nncf config 40 | h, w = cfg.data.height, cfg.data.width 41 | nncf_config_data = get_default_nncf_compression_config(h, w) 42 | 43 | nncf_metainfo = { 44 | 'nncf_compression_enabled': True, 45 | 'nncf_config': nncf_config_data 46 | } 47 | checkpoint['nncf_metainfo'] = nncf_metainfo 48 | res_path = save_checkpoint(checkpoint, args.dst_folder) 49 | 50 | 51 | if __name__ == '__main__': 52 | main() 53 | -------------------------------------------------------------------------------- /tools/auxiliary/compute_mean_std.py: -------------------------------------------------------------------------------- 1 | """ 2 | Compute channel-wise mean and standard deviation of a dataset. 3 | 4 | Usage: 5 | $ python compute_mean_std.py DATASET_ROOT DATASET_KEY 6 | 7 | - The first argument points to the root path where you put the datasets. 8 | - The second argument means the specific dataset key. 9 | 10 | For instance, your datasets are put under $DATA and you wanna 11 | compute the statistics of Market1501, do 12 | $ python compute_mean_std.py $DATA market1501 13 | """ 14 | 15 | import argparse 16 | 17 | import torchreid 18 | 19 | 20 | def main(): 21 | parser = argparse.ArgumentParser() 22 | parser.add_argument('root', type=str) 23 | parser.add_argument('sources', type=str) 24 | args = parser.parse_args() 25 | 26 | datamanager = torchreid.data.ImageDataManager( 27 | root=args.root, 28 | sources=args.sources, 29 | targets=None, 30 | height=256, 31 | width=128, 32 | batch_size_train=100, 33 | batch_size_test=100, 34 | transforms=None, 35 | norm_mean=[0., 0., 0.], 36 | norm_std=[1., 1., 1.], 37 | train_sampler='SequentialSampler' 38 | ) 39 | train_loader = datamanager.train_loader 40 | 41 | print('Computing mean and std ...') 42 | mean = 0. 43 | std = 0. 44 | n_samples = 0. 45 | for data in train_loader: 46 | data = data[0] 47 | batch_size = data.size(0) 48 | data = data.view(batch_size, data.size(1), -1) 49 | mean += data.mean(2).sum(0) 50 | std += data.std(2).sum(0) 51 | n_samples += batch_size 52 | 53 | mean /= n_samples 54 | std /= n_samples 55 | print('Mean: {}'.format(mean)) 56 | print('Std: {}'.format(std)) 57 | 58 | 59 | if __name__ == '__main__': 60 | main() 61 | -------------------------------------------------------------------------------- /tools/auxiliary/explore_weighs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2020-2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import os.path as osp 18 | from argparse import REMAINDER, ArgumentDefaultsHelpFormatter, ArgumentParser 19 | 20 | import numpy as np 21 | import torch 22 | import torch.nn as nn 23 | from scripts.default_config import get_default_config, model_kwargs, merge_from_files_with_base 24 | 25 | import torchreid 26 | from torchreid.utils import load_pretrained_weights 27 | 28 | 29 | def collect_conv_layers(model): 30 | conv_layers = [] 31 | for name, m in model.named_modules(): 32 | if isinstance(m, nn.Conv2d): 33 | weight = m.weight.detach().cpu().numpy() 34 | bias = m.bias.detach().cpu().numpy() if m.bias is not None else 0.0 35 | shape = weight.shape 36 | 37 | assert len(shape) == 4 38 | if shape[2] == shape[3] == 1: 39 | kernel_type = '1x1' 40 | elif shape[1] == 1: 41 | kernel_type = 'dw' 42 | else: 43 | kernel_type = 'reg' 44 | 45 | conv_layers.append(dict( 46 | name=name, 47 | type=kernel_type, 48 | weight=weight, 49 | bias=bias, 50 | updated=False, 51 | )) 52 | elif isinstance(m, nn.BatchNorm2d): 53 | assert len(conv_layers) > 0 54 | 55 | last_conv = conv_layers[-1] 56 | assert not last_conv['updated'] 57 | 58 | alpha = m.weight.detach().cpu().numpy() 59 | beta = m.bias.detach().cpu().numpy() 60 | running_mean = m.running_mean.detach().cpu().numpy() 61 | running_var = m.running_var.detach().cpu().numpy() 62 | 63 | scales = (alpha / np.sqrt(running_var + 1e-5)).reshape([-1, 1, 1, 1]) 64 | last_conv['weight'] = scales * last_conv['weight'] 65 | last_conv['bias'] = scales * (last_conv['bias'] - running_mean) + beta 66 | last_conv['updated'] = True 67 | 68 | return conv_layers 69 | 70 | 71 | def show_stat(conv_layers, max_scale=5.0, max_similarity=0.5, sim_percentile=95): 72 | invalid_weight_scales = [] 73 | invalid_bias_scales = [] 74 | invalid_sim = [] 75 | for conv in conv_layers: 76 | name = conv['name'] 77 | weights = conv['weight'] 78 | bias = conv['bias'] 79 | kernel_type = conv['type'] 80 | if conv['updated']: 81 | kernel_type += ', fused' 82 | 83 | num_filters = weights.shape[0] 84 | filters = weights.reshape([num_filters, -1]) 85 | 86 | norms = np.sqrt(np.sum(np.square(filters), axis=-1)) 87 | min_norm, max_norm = np.min(norms), np.max(norms) 88 | median_norm = np.median(norms) 89 | scale = max_norm / min_norm 90 | 91 | if num_filters <= filters.shape[1] and 'gate.fc' not in name: 92 | norm_filters = filters / norms.reshape([-1, 1]) 93 | similarities = np.matmul(norm_filters, np.transpose(norm_filters)) 94 | 95 | similarities = np.abs(similarities[np.triu_indices(similarities.shape[0], k=1)]) 96 | 97 | num_invalid = np.sum(similarities > max_similarity) 98 | num_total = len(similarities) 99 | if num_invalid > 0: 100 | sim = np.percentile(similarities, sim_percentile) 101 | invalid_sim.append((name, kernel_type, sim, num_invalid, num_total, num_filters)) 102 | 103 | scales = max_norm / norms 104 | num_invalid = np.sum(scales > max_scale) 105 | if num_invalid > 0 or median_norm < 0.1: 106 | invalid_weight_scales.append((name, kernel_type, median_norm, scale, num_invalid, num_filters)) 107 | 108 | bias_scores = np.abs(bias) 109 | bias_score = np.percentile(bias_scores, 95) 110 | if bias_score > 1.0: 111 | invalid_bias_scales.append((name, kernel_type, bias_score)) 112 | 113 | if len(invalid_weight_scales) > 0: 114 | print('\nFound {} layers with invalid weight norm fraction (max/cur > {}):' 115 | .format(len(invalid_weight_scales), max_scale)) 116 | for name, kernel_type, median_norm, scale, num_invalid, num_filters in invalid_weight_scales: 117 | print(' - {} ({}): {:.3f} (median={:.3f} invalid: {} / {})' 118 | .format(name, kernel_type, scale, median_norm, num_invalid, num_filters)) 119 | else: 120 | print('\nThere are no layers with invalid weight norm.') 121 | 122 | if len(invalid_bias_scales) > 0: 123 | print('\nFound {} layers with invalid bias max value (max> {}):' 124 | .format(len(invalid_bias_scales), 1.0)) 125 | for name, kernel_type, scale in invalid_bias_scales: 126 | print(' - {} ({}): {:.3f}' 127 | .format(name, kernel_type, scale)) 128 | else: 129 | print('\nThere are no layers with invalid bias.') 130 | 131 | if len(invalid_sim) > 0: 132 | print('\nFound {} layers with invalid similarity (value > {}):' 133 | .format(len(invalid_sim), max_similarity)) 134 | for name, kernel_type, sim, num_invalid, num_total, num_filters in invalid_sim: 135 | print(' - {} ({}): {:.3f} (invalid: {} / {} size={})' 136 | .format(name, kernel_type, sim, num_invalid, num_total, num_filters)) 137 | else: 138 | print('\nThere are no layers with invalid similarity.') 139 | 140 | 141 | def main(): 142 | parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter) 143 | parser.add_argument('--config', '-c', type=str, required=True) 144 | parser.add_argument('opts', default=None, nargs=REMAINDER) 145 | args = parser.parse_args() 146 | 147 | assert osp.exists(args.config) 148 | 149 | cfg = get_default_config() 150 | cfg.use_gpu = torch.cuda.is_available() 151 | merge_from_files_with_base(cfg, args.config) 152 | cfg.merge_from_list(args.opts) 153 | cfg.freeze() 154 | 155 | model = torchreid.models.build_model(**model_kwargs(cfg, [0, 0])) 156 | load_pretrained_weights(model, cfg.model.load_weights) 157 | 158 | conv_layers = collect_conv_layers(model) 159 | show_stat(conv_layers) 160 | 161 | 162 | if __name__ == '__main__': 163 | main() 164 | -------------------------------------------------------------------------------- /tools/auxiliary/parse_output.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Copyright (C) 2020-2021 Intel Corporation 4 | # SPDX-License-Identifier: Apache-2.0 5 | # 6 | 7 | exp_folder=$1 8 | test_file_path=${exp_folder}/combine_all.txt 9 | for dir in ${exp_folder}/* # list directories in the form "/tmp/dirname/" 10 | do 11 | shopt -s extglob # enable +(...) glob syntax 12 | result=${dir%%+(/)} # trim however many trailing slashes exist 13 | result=${result##*/} # remove everything before the last / that still remains 14 | printf '%s\n' "$result" >> $test_file_path 15 | cat ${dir}/train.log* | grep 'mAP:' >> $test_file_path 16 | cat ${dir}/train.log* | grep 'Rank-1 :' >> $test_file_path 17 | cat ${dir}/train.log* | grep 'Rank-5 :' >> $test_file_path 18 | done 19 | -------------------------------------------------------------------------------- /tools/auxiliary/visualize_actmap.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | """Visualizes CNN activation maps to see where the CNN focuses on to extract features. 9 | 10 | Reference: 11 | - Zagoruyko and Komodakis. Paying more attention to attention: Improving the 12 | performance of convolutional neural networks via attention transfer. ICLR, 2017 13 | - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019. 14 | """ 15 | 16 | import os.path as osp 17 | from argparse import REMAINDER, ArgumentDefaultsHelpFormatter, ArgumentParser 18 | 19 | import cv2 20 | import numpy as np 21 | import torch 22 | import torch.nn.functional as F 23 | from scripts.default_config import (get_default_config, imagedata_kwargs, 24 | model_kwargs, merge_from_files_with_base) 25 | 26 | import torchreid 27 | from torchreid.data.datasets import init_image_dataset 28 | from torchreid.data.transforms import build_transforms 29 | from torchreid.utils import (check_isfile, load_pretrained_weights, 30 | mkdir_if_missing) 31 | 32 | IMAGENET_MEAN = [0.485, 0.456, 0.406] 33 | IMAGENET_STD = [0.229, 0.224, 0.225] 34 | GRID_SPACING = 10 35 | 36 | 37 | def reset_config(cfg, args): 38 | if args.root: 39 | cfg.data.root = args.root 40 | 41 | 42 | def build_dataset(mode='gallery', targets=None, height=192, width=256, 43 | transforms=None, norm_mean=None, norm_std=None, **kwargs): 44 | _, transform_test = build_transforms( 45 | height, 46 | width, 47 | transforms=transforms, 48 | norm_mean=norm_mean, 49 | norm_std=norm_std 50 | ) 51 | 52 | main_target_name = targets[0] 53 | dataset = init_image_dataset( 54 | main_target_name, 55 | transform=transform_test, 56 | mode=mode, 57 | verbose=False, 58 | **kwargs 59 | ) 60 | 61 | return dataset 62 | 63 | 64 | def build_data_loader(dataset, use_gpu=True, batch_size=100): 65 | data_loader = torch.utils.data.DataLoader( 66 | dataset, 67 | batch_size=batch_size, 68 | shuffle=False, 69 | num_workers=1, 70 | pin_memory=use_gpu, 71 | drop_last=False 72 | ) 73 | 74 | return data_loader 75 | 76 | 77 | def prepare_data(cfg, mode='query'): 78 | data_config = imagedata_kwargs(cfg) 79 | dataset = build_dataset(mode=mode, **data_config) 80 | data_loader = build_data_loader(dataset, use_gpu=cfg.use_gpu) 81 | 82 | pids = dataset.num_train_pids 83 | keys = sorted(pids.keys()) 84 | pids = [pids[key] for key in keys] 85 | 86 | return data_loader, pids 87 | 88 | 89 | @torch.no_grad() 90 | def visualize_activation_map(model, data_loader, save_dir, width, height, use_gpu, img_mean=None, img_std=None): 91 | if img_mean is None or img_std is None: 92 | img_mean = IMAGENET_MEAN 93 | img_std = IMAGENET_STD 94 | 95 | model.eval() 96 | 97 | # original images and activation maps are saved individually 98 | actmap_dir = osp.join(save_dir, 'actmap') 99 | mkdir_if_missing(actmap_dir) 100 | 101 | for batch_idx, data in enumerate(data_loader): 102 | imgs, paths = data[0], data[3] 103 | if use_gpu: 104 | imgs = imgs.cuda() 105 | 106 | try: 107 | outputs = model(imgs, return_featuremaps=True) 108 | except TypeError: 109 | raise TypeError( 110 | 'forward() got unexpected keyword argument "return_featuremaps". ' 111 | 'Please add return_featuremaps as an input argument to forward(). When ' 112 | 'return_featuremaps=True, return feature maps only.' 113 | ) 114 | 115 | if outputs.dim() != 4: 116 | raise ValueError( 117 | 'The model output is supposed to have ' 118 | 'shape of (b, c, h, w), i.e. 4 dimensions, but got {} dimensions. ' 119 | 'Please make sure you set the model output at eval mode ' 120 | 'to be the last convolutional feature maps'.format( 121 | outputs.dim() 122 | ) 123 | ) 124 | 125 | # compute activation maps 126 | outputs = outputs.mean(dim=1) 127 | b, h, w = outputs.size() 128 | outputs = outputs.view(b, h * w) 129 | outputs = F.softmax(outputs, dim=1) 130 | # outputs = F.normalize(outputs, p=2, dim=1) 131 | outputs = outputs.view(b, h, w) 132 | 133 | if use_gpu: 134 | imgs, outputs = imgs.cpu(), outputs.cpu() 135 | 136 | for j in range(outputs.size(0)): 137 | # get image name 138 | path = paths[j] 139 | imname = osp.basename(osp.splitext(path)[0]) 140 | 141 | # RGB image 142 | img = imgs[j, ...] 143 | for t, m, s in zip(img, img_mean, img_std): 144 | t.mul_(s).add_(m).clamp_(0, 1) 145 | img_np = np.uint8(np.floor(img.numpy() * 255)) 146 | img_np = img_np.transpose((1, 2, 0)) # (c, h, w) -> (h, w, c) 147 | 148 | # activation map 149 | am = outputs[j, ...].numpy() 150 | am = cv2.resize(am, (width, height)) 151 | am = 255 * (am - np.min(am)) / (np.max(am) - np.min(am) + 1e-12) 152 | am = np.uint8(np.floor(am)) 153 | am = cv2.applyColorMap(am, cv2.COLORMAP_JET) 154 | 155 | # overlapped 156 | overlapped = img_np * 0.3 + am * 0.7 157 | overlapped[overlapped > 255] = 255 158 | overlapped = overlapped.astype(np.uint8) 159 | 160 | # save images in a single figure (add white spacing between images) 161 | # from left to right: original image, activation map, overlapped image 162 | grid_img = 255 * np.ones((height, 3 * width + 2 * GRID_SPACING, 3), dtype=np.uint8) 163 | grid_img[:, :width, :] = img_np[:, :, ::-1] 164 | grid_img[:, width + GRID_SPACING:2 * width + GRID_SPACING, :] = am 165 | grid_img[:, 2 * width + 2 * GRID_SPACING:, :] = overlapped 166 | cv2.imwrite(osp.join(actmap_dir, imname + '.jpg'), grid_img) 167 | 168 | if (batch_idx + 1) % 10 == 0: 169 | print('- done batch {}/{}'.format(batch_idx + 1, len(data_loader))) 170 | 171 | 172 | def main(): 173 | parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter) 174 | parser.add_argument('--config-file', '-c', type=str, required=True) 175 | parser.add_argument('--root', '-r', type=str, required=True) 176 | parser.add_argument('--save-dir', type=str, default='log') 177 | parser.add_argument('opts', default=None, nargs=REMAINDER) 178 | args = parser.parse_args() 179 | 180 | assert osp.exists(args.config_file) 181 | assert osp.exists(args.root) 182 | 183 | cfg = get_default_config() 184 | cfg.use_gpu = torch.cuda.is_available() 185 | if args.config_file: 186 | merge_from_files_with_base(cfg, args.config_file) 187 | reset_config(cfg, args) 188 | cfg.merge_from_list(args.opts) 189 | 190 | if cfg.use_gpu: 191 | torch.backends.cudnn.benchmark = True 192 | 193 | data_loader, num_pids = prepare_data(cfg, mode='gallery') 194 | 195 | print('Building model: {}'.format(cfg.model.name)) 196 | model = torchreid.models.build_model(**model_kwargs(cfg, num_pids)) 197 | 198 | if cfg.model.load_weights and check_isfile(cfg.model.load_weights): 199 | load_pretrained_weights(model, cfg.model.load_weights) 200 | 201 | if cfg.use_gpu: 202 | model = model.cuda() 203 | 204 | visualize_activation_map(model, data_loader, args.save_dir, cfg.data.width, cfg.data.height, cfg.use_gpu) 205 | 206 | 207 | if __name__ == '__main__': 208 | main() 209 | -------------------------------------------------------------------------------- /tools/check_classes_mlc.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | 5 | 6 | def main(): 7 | parser = argparse.ArgumentParser(description='Sample showcasing the new API') 8 | parser.add_argument('data_path', help='path to the dataset root') 9 | parser.add_argument('--train_annotation', 10 | help='path to the pre-trained model weights', 11 | default='train.json') 12 | parser.add_argument('--val_annotation', 13 | help='path to the pre-trained model weights', 14 | default='val.json') 15 | parser.add_argument('--clean', action='store_true') 16 | args = parser.parse_args() 17 | 18 | all_classes = [] 19 | root = args.data_path 20 | subsets=[args.val_annotation, args.train_annotation] 21 | 22 | for subset in subsets: 23 | with open(os.path.join(root, subset)) as f: 24 | anno = json.load(f) 25 | all_classes.append(anno['classes']) 26 | assert len(set(anno['classes'])) == len(anno['classes']) 27 | print(f'Source {subset} len: {len(anno["images"])}') 28 | 29 | if args.clean: 30 | clean_anno = [] 31 | for img_info in anno['images']: 32 | if os.path.isfile(os.path.join(root, img_info[0])): 33 | clean_anno.append(img_info) 34 | else: 35 | print(f'Missing image: {img_info}') 36 | 37 | if len(anno['images']) != len(clean_anno): 38 | anno['images'] = clean_anno 39 | with open(os.path.join(root, subset.split('.')[-2] + '_clean.json', 'w')) as of: 40 | json.dump(anno, of) 41 | 42 | if 'train' in subset: 43 | total_unique_labels = 0 44 | for record in anno['images']: 45 | labels = set(record[-1]) 46 | total_unique_labels += len(labels) 47 | 48 | n_classes = len(anno['classes']) 49 | n_train_images = len(anno['images']) 50 | avg_labels = total_unique_labels / n_train_images 51 | 52 | print(f'Train images {n_train_images}\n' 53 | f'Classes {n_classes}\n' 54 | f'Avg labels per train image {avg_labels:.2f}') 55 | 56 | for i, clssA in enumerate(all_classes): 57 | for j, clssB in enumerate(all_classes): 58 | if j > i: 59 | assert set(clssA) == set(clssB) 60 | 61 | 62 | if __name__ == '__main__': 63 | main() 64 | -------------------------------------------------------------------------------- /tools/convert_to_onnx.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2020-2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import argparse 18 | import os 19 | import warnings 20 | 21 | import torch 22 | 23 | from scripts.default_config import get_default_config, model_kwargs, merge_from_files_with_base 24 | from scripts.script_utils import patch_InplaceAbn_forward 25 | from torchreid.apis.export import export_onnx, export_ir 26 | from torchreid.integration.nncf.compression import get_compression_hyperparams 27 | from torchreid.models import build_model 28 | from torchreid.utils import load_checkpoint, load_pretrained_weights 29 | from torchreid.integration.nncf.compression_script_utils import (make_nncf_changes_in_eval, 30 | make_nncf_changes_in_config) 31 | 32 | 33 | def parse_num_classes(source_datasets, classification=False, num_classes=None, snap_path=None): 34 | if classification: 35 | if snap_path: 36 | chkpt = load_checkpoint(snap_path) 37 | num_classes_from_snap = chkpt['num_classes'] if 'num_classes' in chkpt else None 38 | 39 | if isinstance(num_classes_from_snap, int): 40 | num_classes_from_snap = [num_classes_from_snap] 41 | 42 | if num_classes is None: 43 | num_classes = num_classes_from_snap 44 | else: 45 | print('Warning: number of classes in model was overriden via command line') 46 | assert num_classes is not None and len(num_classes) > 0 47 | 48 | if num_classes is not None and len(num_classes) > 0: 49 | return num_classes 50 | 51 | num_clustered = 0 52 | num_rest = 0 53 | for src in source_datasets: 54 | if isinstance(src, (tuple, list)): 55 | num_clustered += 1 56 | else: 57 | num_rest += 1 58 | 59 | total_num_sources = num_clustered + int(num_rest > 0) 60 | assert total_num_sources > 0 61 | 62 | return [0] * total_num_sources # dummy number of classes 63 | 64 | 65 | def reset_config(cfg): 66 | cfg.model.download_weights = False 67 | 68 | 69 | def main(): 70 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) 71 | parser.add_argument('--config-file', type=str, default='', required=True, 72 | help='Path to config file') 73 | parser.add_argument('--output-name', type=str, default='model', 74 | help='Path to save ONNX model') 75 | parser.add_argument('--num-classes', type=int, nargs='+', default=None) 76 | parser.add_argument('--verbose', action='store_true', 77 | help='Verbose mode for onnx.export') 78 | parser.add_argument('--disable-dyn-axes', default=False, action='store_true') 79 | parser.add_argument('--export_ir', action='store_true') 80 | parser.add_argument('opts', default=None, nargs=argparse.REMAINDER, 81 | help='Modify config options using the command-line') 82 | args = parser.parse_args() 83 | 84 | cfg = get_default_config() 85 | cfg.use_gpu = torch.cuda.is_available() 86 | if args.config_file: 87 | merge_from_files_with_base(cfg, args.config_file) 88 | reset_config(cfg) 89 | cfg.merge_from_list(args.opts) 90 | 91 | compression_hyperparams = get_compression_hyperparams(cfg.model.load_weights) 92 | is_nncf_used = compression_hyperparams['enable_quantization'] or compression_hyperparams['enable_pruning'] 93 | if is_nncf_used: 94 | print('Using NNCF -- making NNCF changes in config') 95 | cfg = make_nncf_changes_in_config(cfg, 96 | compression_hyperparams['enable_quantization'], 97 | compression_hyperparams['enable_pruning'], 98 | args.opts) 99 | cfg.train.mix_precision = False 100 | cfg.freeze() 101 | num_classes = parse_num_classes(source_datasets=cfg.data.sources, 102 | classification=cfg.model.type in ('classification', 'multilabel'), 103 | num_classes=args.num_classes, 104 | snap_path=cfg.model.load_weights) 105 | model = build_model(**model_kwargs(cfg, num_classes)) 106 | if cfg.model.load_weights: 107 | load_pretrained_weights(model, cfg.model.load_weights) 108 | else: 109 | warnings.warn("No weights are passed through 'load_weights' parameter! " 110 | "The model will be converted with random or pretrained weights", category=RuntimeWarning) 111 | if 'tresnet' in cfg.model.name: 112 | patch_InplaceAbn_forward() 113 | if is_nncf_used: 114 | print('Begin making NNCF changes in model') 115 | model = make_nncf_changes_in_eval(model, cfg) 116 | print('End making NNCF changes in model') 117 | onnx_file_path = export_onnx(model=model.eval(), 118 | cfg=cfg, 119 | output_file_path=args.output_name, 120 | disable_dyn_axes=args.disable_dyn_axes, 121 | verbose=args.verbose, 122 | opset=cfg.model.export_onnx_opset, 123 | extra_check=True) 124 | if args.export_ir: 125 | input_shape = [1, 3, cfg.data.height, cfg.data.width] 126 | export_ir(onnx_model_path=onnx_file_path, 127 | norm_mean=cfg.data.norm_mean, 128 | norm_std=cfg.data.norm_std, 129 | input_shape=input_shape, 130 | optimized_model_dir=os.path.dirname(os.path.abspath(onnx_file_path)), 131 | data_type='FP32') 132 | 133 | 134 | if __name__ == '__main__': 135 | main() 136 | -------------------------------------------------------------------------------- /tools/eval.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import os.path as osp 18 | import time 19 | import sys 20 | 21 | import torch 22 | from scripts.default_config import get_default_config, model_kwargs, imagedata_kwargs, merge_from_files_with_base 23 | from scripts.script_utils import reset_config, build_base_argparser, check_classification_classes 24 | 25 | import torchreid 26 | from torchreid.integration.nncf.compression import get_compression_hyperparams 27 | from torchreid.ops import DataParallel 28 | from torchreid.utils import (Logger, set_random_seed, load_pretrained_weights,) 29 | from torchreid.engine import build_engine 30 | from torchreid.integration.nncf.compression_script_utils import (make_nncf_changes_in_eval, 31 | make_nncf_changes_in_config) 32 | from torchreid.metrics.classification import score_extraction, tune_multilabel_thresholds 33 | 34 | 35 | def main(): 36 | parser = build_base_argparser() 37 | args = parser.parse_args() 38 | 39 | cfg = get_default_config() 40 | cfg.use_gpu = torch.cuda.is_available() and args.gpu_num > 0 41 | if args.config_file: 42 | merge_from_files_with_base(cfg, args.config_file) 43 | reset_config(cfg, args) 44 | cfg.merge_from_list(args.opts) 45 | 46 | is_ie_model = cfg.model.load_weights.endswith('.xml') 47 | if not is_ie_model: 48 | compression_hyperparams = get_compression_hyperparams(cfg.model.load_weights) 49 | is_nncf_used = compression_hyperparams['enable_quantization'] or compression_hyperparams['enable_pruning'] 50 | 51 | if is_nncf_used: 52 | print('Using NNCF -- making NNCF changes in config') 53 | cfg = make_nncf_changes_in_config(cfg, 54 | compression_hyperparams['enable_quantization'], 55 | compression_hyperparams['enable_pruning'], 56 | args.opts) 57 | else: 58 | is_nncf_used = False 59 | 60 | set_random_seed(cfg.train.seed) 61 | 62 | log_name = 'test.log' + time.strftime('-%Y-%m-%d-%H-%M-%S') 63 | sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name)) 64 | datamanager = torchreid.data.ImageDataManager(filter_classes=args.classes, **imagedata_kwargs(cfg)) 65 | num_classes = len(datamanager.test_loader.dataset.classes) 66 | cfg.train.ema.enable = False 67 | if not is_ie_model: 68 | model = torchreid.models.build_model(**model_kwargs(cfg, num_classes)) 69 | load_pretrained_weights(model, cfg.model.load_weights) 70 | if is_nncf_used: 71 | print('Begin making NNCF changes in model') 72 | model = make_nncf_changes_in_eval(model, cfg) 73 | print('End making NNCF changes in model') 74 | if cfg.use_gpu: 75 | num_devices = min(torch.cuda.device_count(), args.gpu_num) 76 | main_device_ids = list(range(num_devices)) 77 | model = DataParallel(model, device_ids=main_device_ids, output_device=0).cuda(main_device_ids[0]) 78 | else: 79 | from torchreid.utils.ie_tools import VectorCNN 80 | from openvino.inference_engine import IECore 81 | cfg.test.batch_size = 1 82 | model = VectorCNN(IECore(), cfg.model.load_weights, 'CPU', switch_rb=True, **model_kwargs(cfg, num_classes)) 83 | for _, dataloader in datamanager.test_loader.items(): 84 | dataloader['query'].dataset.transform.transforms = \ 85 | dataloader['query'].dataset.transform.transforms[:-2] 86 | 87 | if cfg.model.type == 'classification': 88 | check_classification_classes(model, datamanager, args.classes, test_only=True) 89 | 90 | engine = build_engine(cfg=cfg, datamanager=datamanager, model=model, optimizer=None, scheduler=None) 91 | extra_args = {} 92 | if cfg.test.estimate_multilabel_thresholds and cfg.model.type == 'multilabel': 93 | print('Estimation optimal thresholds ...') 94 | scores, labels = score_extraction(datamanager.train_loader, model, cfg.use_gpu) 95 | thresholds = tune_multilabel_thresholds(scores, labels) 96 | extra_args['pos_thresholds'] = thresholds 97 | engine.test(0, topk=(1, 5, 10, 20), test_only=True, **extra_args) 98 | 99 | 100 | if __name__ == '__main__': 101 | main() 102 | -------------------------------------------------------------------------------- /tools/get_flops.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2020-2021 Intel Corporation 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | 5 | import argparse 6 | import json 7 | 8 | import torch 9 | from scripts.default_config import (get_default_config, imagedata_kwargs, 10 | model_kwargs, merge_from_files_with_base) 11 | 12 | import torchreid 13 | from torchreid.utils import set_random_seed 14 | from ptflops import get_model_complexity_info 15 | from fvcore.nn import FlopCountAnalysis, flop_count_table 16 | 17 | def build_datamanager(cfg, classification_classes_filter=None): 18 | return torchreid.data.ImageDataManager(filter_classes=classification_classes_filter, **imagedata_kwargs(cfg)) 19 | 20 | 21 | def reset_config(cfg, args): 22 | if args.root: 23 | cfg.data.root = args.root 24 | if args.custom_roots: 25 | cfg.custom_datasets.roots = args.custom_roots 26 | if args.custom_types: 27 | cfg.custom_datasets.types = args.custom_types 28 | if args.custom_names: 29 | cfg.custom_datasets.names = args.custom_names 30 | 31 | 32 | def main(): 33 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) 34 | parser.add_argument('--config-file', type=str, default='', required=True, 35 | help='path to config file') 36 | parser.add_argument('--custom-roots', type=str, nargs='+', 37 | help='types or paths to annotation of custom datasets (delimited by space)') 38 | parser.add_argument('--custom-types', type=str, nargs='+', 39 | help='path of custom datasets (delimited by space)') 40 | parser.add_argument('--custom-names', type=str, nargs='+', 41 | help='names of custom datasets (delimited by space)') 42 | parser.add_argument('--root', type=str, default='', help='path to data root') 43 | parser.add_argument('--classes', type=str, nargs='+', 44 | help='name of classes in classification dataset') 45 | parser.add_argument('--out') 46 | parser.add_argument('--fvcore', action='store_true', 47 | help='option to use fvcore tool from Meta Platforms for the flops counting') 48 | parser.add_argument('opts', default=None, nargs=argparse.REMAINDER, 49 | help='Modify config options using the command-line') 50 | args = parser.parse_args() 51 | 52 | cfg = get_default_config() 53 | cfg.use_gpu = torch.cuda.is_available() 54 | if args.config_file: 55 | merge_from_files_with_base(cfg, args.config_file) 56 | reset_config(cfg, args) 57 | cfg.merge_from_list(args.opts) 58 | set_random_seed(cfg.train.seed) 59 | 60 | print(f'Show configuration\n{cfg}\n') 61 | 62 | if cfg.use_gpu: 63 | torch.backends.cudnn.benchmark = True 64 | 65 | datamanager = build_datamanager(cfg, args.classes) 66 | num_train_classes = datamanager.num_train_ids 67 | 68 | print(f'Building main model: {cfg.model.name}') 69 | model = torchreid.models.build_model(**model_kwargs(cfg, num_train_classes)) 70 | macs, num_params = get_model_complexity_info(model, (3, cfg.data.height, cfg.data.width), 71 | as_strings=False, verbose=False, print_per_layer_stat=False) 72 | print(f'Main model complexity: M params={num_params / 10**6:,} G flops={macs * 2 / 10**9:,}') 73 | 74 | if args.fvcore: 75 | input_ = torch.rand((1, 3, cfg.data.height, cfg.data.width), dtype=next(model.parameters()).dtype, 76 | device=next(model.parameters()).device) 77 | flops = FlopCountAnalysis(model, input_) 78 | print(f"Main model complexity by fvcore: {flops.total()}") 79 | print(flop_count_table(flops)) 80 | 81 | if args.out: 82 | out = [] 83 | out.append({'key': 'size', 'display_name': 'Size', 'value': num_params / 10**6, 'unit': 'Mp'}) 84 | out.append({'key': 'complexity', 'display_name': 'Complexity', 'value': 2 * macs / 10**9, 85 | 'unit': 'GFLOPs'}) 86 | print('dump to' + args.out) 87 | with open(args.out, 'w') as write_file: 88 | json.dump(out, write_file, indent=4) 89 | 90 | if __name__ == '__main__': 91 | main() 92 | -------------------------------------------------------------------------------- /tools/main.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2020-2021 Intel Corporation 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | 5 | import os.path as osp 6 | import sys 7 | import time 8 | 9 | from torch.utils.tensorboard import SummaryWriter 10 | import torch 11 | from ptflops import get_model_complexity_info 12 | 13 | from scripts.default_config import (get_default_config, 14 | lr_scheduler_kwargs, model_kwargs, 15 | optimizer_kwargs, 16 | merge_from_files_with_base) 17 | from scripts.script_utils import (build_base_argparser, reset_config, 18 | check_classification_classes, 19 | build_datamanager, 20 | put_main_model_on_the_device) 21 | 22 | import torchreid 23 | from torchreid.apis.training import run_lr_finder, run_training 24 | from torchreid.utils import (Logger, check_isfile, resume_from_checkpoint, 25 | set_random_seed, load_pretrained_weights, get_git_revision) 26 | from torchreid.integration.nncf.compression_script_utils import (make_nncf_changes_in_config, 27 | make_nncf_changes_in_training) 28 | 29 | 30 | def main(): 31 | parser = build_base_argparser() 32 | parser.add_argument('-e', '--auxiliary-models-cfg', type=str, nargs='*', default='', 33 | help='path to extra config files') 34 | parser.add_argument('--split-models', action='store_true', 35 | help='whether to split models on own gpu') 36 | parser.add_argument('--enable_quantization', action='store_true', 37 | help='Enable NNCF quantization algorithm') 38 | parser.add_argument('--enable_pruning', action='store_true', 39 | help='Enable NNCF pruning algorithm') 40 | parser.add_argument('--aux-config-opts', nargs='+', default=None, 41 | help='Modify aux config options using the command-line') 42 | args = parser.parse_args() 43 | 44 | cfg = get_default_config() 45 | cfg.use_gpu = torch.cuda.is_available() and args.gpu_num > 0 46 | if args.config_file: 47 | merge_from_files_with_base(cfg, args.config_file) 48 | reset_config(cfg, args) 49 | cfg.merge_from_list(args.opts) 50 | 51 | is_nncf_used = args.enable_quantization or args.enable_pruning 52 | if is_nncf_used: 53 | print('Using NNCF -- making NNCF changes in config') 54 | cfg = make_nncf_changes_in_config(cfg, 55 | args.enable_quantization, 56 | args.enable_pruning, 57 | args.opts) 58 | 59 | set_random_seed(cfg.train.seed, cfg.train.deterministic) 60 | 61 | log_name = 'test.log' if cfg.test.evaluate else 'train.log' 62 | log_name += time.strftime('-%Y-%m-%d-%H-%M-%S') 63 | sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name)) 64 | sha_commit, branch_name = get_git_revision() 65 | print(f'HEAD is: {branch_name}') 66 | print(f'commit SHA is: {sha_commit}\n') 67 | print(f'Show configuration\n{cfg}\n') 68 | 69 | if cfg.use_gpu: 70 | torch.backends.cudnn.benchmark = True 71 | 72 | num_aux_models = len(cfg.mutual_learning.aux_configs) 73 | datamanager = build_datamanager(cfg, args.classes) 74 | num_train_classes = datamanager.num_train_ids 75 | 76 | print(f'Building main model: {cfg.model.name}') 77 | model = torchreid.models.build_model(**model_kwargs(cfg, num_train_classes)) 78 | macs, num_params = get_model_complexity_info(model, (3, cfg.data.height, cfg.data.width), 79 | as_strings=False, verbose=False, print_per_layer_stat=False) 80 | print(f'Main model complexity: params={num_params:,} flops={macs * 2:,}') 81 | 82 | aux_lr = cfg.train.lr # placeholder, needed for aux models, may be filled by nncf part below 83 | if is_nncf_used: 84 | print('Begin making NNCF changes in model') 85 | if cfg.use_gpu: 86 | model.cuda() 87 | 88 | compression_ctrl, model, cfg, aux_lr, nncf_metainfo = \ 89 | make_nncf_changes_in_training(model, cfg, 90 | args.classes, 91 | args.opts) 92 | 93 | should_freeze_aux_models = True 94 | print(f'should_freeze_aux_models = {should_freeze_aux_models}') 95 | print('End making NNCF changes in model') 96 | else: 97 | compression_ctrl = None 98 | should_freeze_aux_models = False 99 | nncf_metainfo = None 100 | # creating optimizer and scheduler -- it should be done after NNCF part, since 101 | # NNCF could change some parameters 102 | optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg)) 103 | 104 | if cfg.lr_finder.enable and not cfg.model.resume: 105 | scheduler = None 106 | else: 107 | scheduler = torchreid.optim.build_lr_scheduler(optimizer=optimizer, 108 | num_iter=datamanager.num_iter, 109 | **lr_scheduler_kwargs(cfg)) 110 | # Loading model (and optimizer and scheduler in case of resuming training). 111 | # Note that if NNCF is used, loading is done inside NNCF part, so loading here is not required. 112 | if cfg.model.resume and check_isfile(cfg.model.resume) and not is_nncf_used: 113 | device_ = 'cuda' if cfg.use_gpu else 'cpu' 114 | cfg.train.start_epoch = resume_from_checkpoint( 115 | cfg.model.resume, model, optimizer=optimizer, scheduler=scheduler, device=device_ 116 | ) 117 | elif cfg.model.load_weights and not is_nncf_used: 118 | load_pretrained_weights(model, cfg.model.load_weights) 119 | 120 | if cfg.model.type == 'classification': 121 | check_classification_classes(model, datamanager, args.classes, test_only=cfg.test.evaluate) 122 | 123 | model, extra_device_ids = put_main_model_on_the_device(model, cfg.use_gpu, args.gpu_num, 124 | num_aux_models, args.split_models) 125 | 126 | if cfg.lr_finder.enable and not cfg.test.evaluate and not cfg.model.resume: 127 | aux_lr, model, optimizer, scheduler = run_lr_finder(cfg, datamanager, model, 128 | optimizer, scheduler, args.classes, 129 | rebuild_model=True, 130 | gpu_num=args.gpu_num, 131 | split_models=args.split_models) 132 | 133 | log_dir = cfg.data.tb_log_dir if cfg.data.tb_log_dir else cfg.data.save_dir 134 | run_training(cfg, datamanager, model, optimizer, scheduler, extra_device_ids, 135 | aux_lr, tb_writer=SummaryWriter(log_dir=log_dir), 136 | aux_config_opts=args.aux_config_opts, 137 | should_freeze_aux_models=should_freeze_aux_models, 138 | nncf_metainfo=nncf_metainfo, 139 | compression_ctrl=compression_ctrl) 140 | 141 | 142 | if __name__ == '__main__': 143 | main() 144 | -------------------------------------------------------------------------------- /tools/run_pylint.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | 4 | from subprocess import run # nosec 5 | 6 | if __name__ == '__main__': 7 | ignored_patterns = [ 8 | 'setup.py', 9 | 'tools/auxiliary/', 10 | '.history', 11 | 'torchreid/models', 12 | 'torchreid/optim/radam.py', 13 | 'torchreid/optim/sam.py', 14 | 'tests/conftest.py', 15 | 'tests/config.py', 16 | 'torchreid/integration/sc/model_wrappers/classification.py', 17 | 'tests/test_ote_training.py' 18 | ] 19 | 20 | to_pylint = [] 21 | wd = os.path.abspath('.') 22 | for root, dirnames, filenames in os.walk(wd): 23 | for filename in filenames: 24 | if filename.endswith('.py'): 25 | full_path = os.path.join(root, filename) 26 | rel_path = os.path.relpath(full_path, wd) 27 | if all(not re.match(pattern, rel_path) for pattern in ignored_patterns): 28 | to_pylint.append(rel_path) 29 | 30 | run(['pylint'] + to_pylint, check=True) 31 | -------------------------------------------------------------------------------- /torchreid/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, print_function 9 | 10 | try: 11 | import nncf 12 | except ImportError: 13 | import warnings 14 | warnings.warn("NNCF was not found. Model optimization options will not be available.") 15 | 16 | from torchreid import data, engine, losses, metrics, models, ops, optim, utils 17 | from .version import __version__ 18 | 19 | __author__ = 'Kaiyang Zhou' 20 | __homepage__ = 'https://kaiyangzhou.github.io/' 21 | __description__ = 'Deep learning person re-identification in PyTorch' 22 | __url__ = 'https://github.com/KaiyangZhou/deep-person-reid' 23 | -------------------------------------------------------------------------------- /torchreid/apis/export.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from subprocess import run, DEVNULL, CalledProcessError # nosec 18 | 19 | import onnx 20 | import torch 21 | 22 | from torch.onnx.symbolic_registry import register_op 23 | 24 | from scripts.script_utils import (group_norm_symbolic, hardsigmoid_symbolic, 25 | relu6_symbolic) 26 | from torchreid.utils import random_image 27 | from torchreid.data.transforms import build_inference_transform 28 | 29 | 30 | def export_onnx(model, cfg, output_file_path='model', disable_dyn_axes=True, 31 | verbose=False, opset=9, extra_check=False, output_names=['output']): 32 | transform = build_inference_transform( 33 | cfg.data.height, 34 | cfg.data.width, 35 | norm_mean=cfg.data.norm_mean, 36 | norm_std=cfg.data.norm_std, 37 | ) 38 | 39 | input_img = random_image(cfg.data.height, cfg.data.width) 40 | device = next(model.parameters()).device 41 | input_blob = transform(input_img).unsqueeze(0).to(device) 42 | 43 | input_names = ['data'] 44 | if not disable_dyn_axes: 45 | dynamic_axes = {'data': {0: 'batch_size', 1: 'channels', 2: 'height', 3: 'width'}} 46 | else: 47 | dynamic_axes = {} 48 | 49 | if not output_file_path.endswith('.onnx'): 50 | output_file_path += '.onnx' 51 | 52 | register_op("group_norm", group_norm_symbolic, "", opset) 53 | register_op("relu6", relu6_symbolic, "", opset) 54 | register_op("hardsigmoid", hardsigmoid_symbolic, "", opset) 55 | 56 | with torch.no_grad(): 57 | torch.onnx.export( 58 | model, 59 | input_blob, 60 | output_file_path, 61 | verbose=verbose, 62 | export_params=True, 63 | input_names=input_names, 64 | output_names=output_names, 65 | dynamic_axes=dynamic_axes, 66 | opset_version=opset, 67 | operator_export_type=torch.onnx.OperatorExportTypes.ONNX, 68 | ) 69 | 70 | if extra_check: 71 | net_from_onnx = onnx.load(output_file_path) 72 | try: 73 | onnx.checker.check_model(net_from_onnx) 74 | print('ONNX check passed.') 75 | except onnx.onnx_cpp2py_export.checker.ValidationError as ex: 76 | print(f'ONNX check failed: {ex}.') 77 | 78 | return output_file_path 79 | 80 | 81 | def export_ir(onnx_model_path, norm_mean=[0,0,0], norm_std=[1,1,1], input_shape=None, 82 | optimized_model_dir='./ir_model', data_type='FP32', pruning_transformation=False): 83 | def get_mo_cmd(): 84 | for mo_cmd in ('mo', 'mo.py'): 85 | try: 86 | run([mo_cmd, '-h'], stdout=DEVNULL, stderr=DEVNULL, shell=False, check=True) 87 | return mo_cmd 88 | except CalledProcessError: 89 | pass 90 | raise RuntimeError('OpenVINO Model Optimizer is not found or configured improperly') 91 | 92 | mean_values = str([s*255 for s in norm_mean]) 93 | scale_values = str([s*255 for s in norm_std]) 94 | 95 | mo_cmd = get_mo_cmd() 96 | 97 | command_line = [mo_cmd, f'--input_model={onnx_model_path}', 98 | f'--mean_values={mean_values}', 99 | f'--scale_values={scale_values}', 100 | f'--output_dir={optimized_model_dir}', 101 | '--data_type', f'{data_type}', 102 | '--reverse_input_channels'] 103 | 104 | if input_shape: 105 | command_line.extend(['--input_shape', f"{input_shape}"]) 106 | 107 | if pruning_transformation: 108 | command_line.extend(['--transform', 'Pruning']) 109 | 110 | run(command_line, shell=False, check=True) 111 | -------------------------------------------------------------------------------- /torchreid/data/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, print_function 2 | 3 | from .datamanager import ImageDataManager 4 | -------------------------------------------------------------------------------- /torchreid/data/datamanager.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, division, print_function 9 | 10 | import numpy as np 11 | import torch 12 | 13 | from torchreid.data.datasets import init_image_dataset 14 | from torchreid.data.sampler import build_train_sampler 15 | from torchreid.data.transforms import build_transforms 16 | from torchreid.utils import worker_init_fn 17 | 18 | 19 | class ImageDataManager(): 20 | r"""Image data manager. 21 | 22 | Args: 23 | root (str): root path to datasets. 24 | height (int, optional): target image height. Default is 256. 25 | width (int, optional): target image width. Default is 128. 26 | transforms (str or list of str, optional): transformations applied to model training. 27 | Default is 'random_flip'. 28 | norm_mean (list or None, optional): data mean. Default is None (use imagenet mean). 29 | norm_std (list or None, optional): data std. Default is None (use imagenet std). 30 | use_gpu (bool, optional): use gpu. Default is True. 31 | batch_size_train (int, optional): number of images in a training batch. Default is 32. 32 | batch_size_test (int, optional): number of images in a test batch. Default is 32. 33 | workers (int, optional): number of workers. Default is 4. 34 | train_sampler (str, optional): sampler. Default is RandomSampler. 35 | correct_batch_size (bool, optional): this heuristic improves multilabel training on small datasets 36 | """ 37 | 38 | def __init__( 39 | self, 40 | root='', 41 | height=224, 42 | width=224, 43 | transforms='random_flip', 44 | norm_mean=None, 45 | norm_std=None, 46 | use_gpu=True, 47 | batch_size_train=32, 48 | batch_size_test=32, 49 | correct_batch_size = False, 50 | workers=4, 51 | train_sampler='RandomSampler', 52 | custom_dataset_roots=[''], 53 | custom_dataset_types=[''], 54 | filter_classes=None, 55 | ): 56 | self.height = height 57 | self.width = width 58 | 59 | self.transform_train, self.transform_test = build_transforms( 60 | self.height, 61 | self.width, 62 | transforms=transforms, 63 | norm_mean=norm_mean, 64 | norm_std=norm_std, 65 | ) 66 | 67 | self.use_gpu = (torch.cuda.is_available() and use_gpu) 68 | 69 | print('=> Loading train dataset') 70 | train_dataset = init_image_dataset( 71 | transform=self.transform_train, 72 | mode='train', 73 | root=root, 74 | custom_dataset_roots=custom_dataset_roots, 75 | custom_dataset_types=custom_dataset_types, 76 | filter_classes=filter_classes, 77 | ) 78 | 79 | if correct_batch_size: 80 | batch_size_train = self.calculate_batch(batch_size_train, len(train_dataset)) 81 | batch_size_train = max(1, min(batch_size_train, len(train_dataset))) 82 | self.train_loader = torch.utils.data.DataLoader( 83 | train_dataset, 84 | sampler=build_train_sampler( 85 | train_dataset.data, 86 | train_sampler, 87 | ), 88 | batch_size=batch_size_train, 89 | shuffle=False, 90 | worker_init_fn=worker_init_fn, 91 | num_workers=workers, 92 | pin_memory=self.use_gpu, 93 | drop_last=True 94 | ) 95 | self.num_iter = len(self.train_loader) 96 | print('=> Loading test dataset') 97 | 98 | # build test loader 99 | test_dataset = init_image_dataset( 100 | transform=self.transform_test, 101 | mode='test', 102 | root=root, 103 | custom_dataset_roots=custom_dataset_roots, 104 | custom_dataset_types=custom_dataset_types, 105 | filter_classes=filter_classes, 106 | ) 107 | self._num_train_ids = train_dataset.num_ids 108 | assert train_dataset.num_ids == test_dataset.num_ids 109 | 110 | self.test_loader = torch.utils.data.DataLoader( 111 | test_dataset, 112 | batch_size=max(min(batch_size_test, len(test_dataset)), 1), 113 | shuffle=False, 114 | num_workers=workers, 115 | worker_init_fn=worker_init_fn, 116 | pin_memory=self.use_gpu, 117 | drop_last=False 118 | ) 119 | 120 | print('\n') 121 | print(' **************** Summary ****************') 122 | print(f' # categories : {self._num_train_ids}') 123 | print(f' # train images : {len(train_dataset)}') 124 | print(f' # test images : {len(test_dataset)}') 125 | print(' *****************************************') 126 | print('\n') 127 | 128 | @staticmethod 129 | def calculate_batch(cur_batch, data_len): 130 | ''' This heuristic improves multilabel training on small datasets ''' 131 | if data_len <= 2500: 132 | return max(int(np.ceil(np.sqrt(data_len) / 2.5)), 6) 133 | return cur_batch 134 | 135 | @property 136 | def num_train_ids(self): 137 | """Returns the number of training categories.""" 138 | return self._num_train_ids 139 | -------------------------------------------------------------------------------- /torchreid/data/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, print_function 9 | from copy import copy 10 | 11 | from .image import (Classification, ClassificationImageFolder, ExternalDatasetWrapper, 12 | MultiLabelClassification, MultiheadClassification) 13 | 14 | __image_datasets = { 15 | 'classification': Classification, 16 | 'classification_image_folder' : ClassificationImageFolder, 17 | 'external_classification_wrapper' : ExternalDatasetWrapper, 18 | 'multilabel_classification': MultiLabelClassification, 19 | 'multihead_classification': MultiheadClassification, 20 | } 21 | 22 | 23 | def init_image_dataset(mode, 24 | custom_dataset_roots=[''], 25 | custom_dataset_types=[''], 26 | **kwargs): 27 | """Initializes an image dataset.""" 28 | 29 | # handle also custom datasets 30 | avai_datasets = list(__image_datasets.keys()) 31 | for data_type in custom_dataset_types: 32 | assert data_type in avai_datasets 33 | new_kwargs = copy(kwargs) 34 | i = 0 if mode == 'train' else 1 35 | if custom_dataset_types[i] == 'external_classification_wrapper': 36 | new_kwargs['data_provider'] = custom_dataset_roots[i] 37 | else: 38 | new_kwargs['root'] = custom_dataset_roots[i] 39 | return __image_datasets[custom_dataset_types[i]](**new_kwargs) 40 | -------------------------------------------------------------------------------- /torchreid/data/datasets/dataset.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, division, print_function 9 | import os.path as osp 10 | 11 | import torch 12 | 13 | from torchreid.utils import read_image 14 | 15 | 16 | class ImageDataset: 17 | """A base class representing ImageDataset. 18 | 19 | All other image datasets should subclass it. 20 | 21 | ``__getitem__`` returns an image given index. 22 | It will return ``img``, ``img_path`` 23 | where ``img`` has shape (channel, height, width). As a result, 24 | data in each batch has shape (batch_size, channel, height, width). 25 | """ 26 | 27 | def __init__(self, 28 | data, 29 | transform=None, 30 | verbose=True, 31 | mixed_cls_heads_info={}, 32 | classes={}, 33 | num_ids=0, 34 | **kwargs): 35 | 36 | self.classes = classes 37 | self.mixed_cls_heads_info = mixed_cls_heads_info 38 | self.data = data 39 | self.transform = transform 40 | self.verbose = verbose 41 | self.num_ids = num_ids 42 | 43 | def __getitem__(self, index): 44 | input_record = self.data[index] 45 | 46 | image = read_image(input_record[0], grayscale=False) 47 | obj_id = input_record[1] 48 | 49 | if isinstance(obj_id, (tuple, list)): # when multi-label classification is available 50 | if len(self.mixed_cls_heads_info): 51 | targets = torch.IntTensor(obj_id) 52 | else: 53 | targets = torch.zeros(self.num_ids) 54 | for obj in obj_id: 55 | targets[obj] = 1 56 | obj_id = targets 57 | 58 | if self.transform is not None: 59 | transformed_image = self.transform(image) 60 | else: 61 | transformed_image = image 62 | 63 | output_record = (transformed_image, obj_id) 64 | 65 | return output_record 66 | 67 | def __len__(self): 68 | return len(self.data) 69 | 70 | @staticmethod 71 | def check_before_run(required_files): 72 | """Checks if required files exist before going deeper. 73 | 74 | Args: 75 | required_files (str or list): string file name(s). 76 | """ 77 | if isinstance(required_files, str): 78 | required_files = [required_files] 79 | 80 | for fpath in required_files: 81 | if not osp.exists(fpath): 82 | raise RuntimeError(f'"{fpath}" is not found') 83 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, print_function 9 | 10 | from .classification import (Classification, ClassificationImageFolder, 11 | ExternalDatasetWrapper, MultiLabelClassification, 12 | MultiheadClassification) 13 | -------------------------------------------------------------------------------- /torchreid/data/sampler.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, division 9 | 10 | from torch.utils.data.sampler import RandomSampler, SequentialSampler 11 | 12 | AVAI_SAMPLERS = ['SequentialSampler', 'RandomSampler'] 13 | 14 | 15 | def build_train_sampler(data_source, train_sampler, **kwargs): 16 | """Builds a training sampler. 17 | """ 18 | assert train_sampler in AVAI_SAMPLERS, \ 19 | f'train_sampler must be one of {AVAI_SAMPLERS}, but got {train_sampler}' 20 | 21 | if train_sampler == 'SequentialSampler': 22 | sampler = SequentialSampler(data_source) 23 | elif train_sampler == 'RandomSampler': 24 | sampler = RandomSampler(data_source) 25 | else: 26 | raise ValueError(f'Unknown sampler: {train_sampler}') 27 | 28 | return sampler 29 | -------------------------------------------------------------------------------- /torchreid/engine/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from __future__ import print_function, absolute_import 3 | 4 | from .engine import Engine, EpochIntervalToValue 5 | from .image import MultilabelEngine, ImageAMSoftmaxEngine, MultiheadEngine 6 | from .builder import build_engine 7 | -------------------------------------------------------------------------------- /torchreid/engine/builder.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2020-2021 Intel Corporation 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | 5 | from .image import (ImageAMSoftmaxEngine, MultilabelEngine, MultiheadEngine) 6 | 7 | NNCF_ENABLED_LOSS = ['softmax', 'am_softmax', 'am_binary', 'asl'] 8 | 9 | def build_engine(cfg, datamanager, model, optimizer, scheduler, 10 | should_freeze_aux_models=False, 11 | nncf_metainfo=None, 12 | compression_ctrl=None, 13 | initial_lr=None): 14 | if should_freeze_aux_models or nncf_metainfo: 15 | if not all(loss_name in NNCF_ENABLED_LOSS for loss_name in cfg.loss.name.split(',')): 16 | raise NotImplementedError('Freezing of aux models or NNCF compression are supported only for ' 17 | 'softmax, am_softmax and am_binary losses for data.type = image') 18 | initial_lr = initial_lr if initial_lr else cfg.train.lr 19 | classification_params = dict( 20 | datamanager=datamanager, 21 | models=model, 22 | optimizers=optimizer, 23 | schedulers=scheduler, 24 | use_gpu=cfg.use_gpu, 25 | save_all_chkpts = cfg.model.save_all_chkpts, 26 | lr_finder = cfg.lr_finder.enable, 27 | train_patience = cfg.train.train_patience, 28 | early_stopping = cfg.train.early_stopping, 29 | lr_decay_factor = cfg.train.lr_decay_factor, 30 | conf_penalty=cfg.loss.softmax.conf_penalty, 31 | label_smooth=cfg.loss.softmax.label_smooth, 32 | aug_type=cfg.loss.softmax.augmentations.aug_type, 33 | aug_prob=cfg.loss.softmax.augmentations.aug_prob, 34 | decay_power=cfg.loss.softmax.augmentations.fmix.decay_power, 35 | alpha=cfg.loss.softmax.augmentations.alpha, 36 | pr_product=cfg.loss.softmax.pr_product, 37 | loss_name=cfg.loss.name, 38 | clip_grad=cfg.train.clip_grad, 39 | m=cfg.loss.softmax.m, 40 | s=cfg.loss.softmax.s, 41 | compute_s=cfg.loss.softmax.compute_s, 42 | margin_type=cfg.loss.softmax.margin_type, 43 | symmetric_ce=cfg.loss.softmax.symmetric_ce, 44 | enable_rsc=cfg.model.self_challenging_cfg.enable, 45 | should_freeze_aux_models=should_freeze_aux_models, 46 | nncf_metainfo=nncf_metainfo, 47 | compression_ctrl=compression_ctrl, 48 | initial_lr=initial_lr, 49 | target_metric=cfg.train.target_metric, 50 | use_ema_decay=cfg.train.ema.enable, 51 | ema_decay=cfg.train.ema.ema_decay, 52 | asl_gamma_neg=cfg.loss.asl.gamma_neg, 53 | asl_gamma_pos=cfg.loss.asl.gamma_pos, 54 | asl_p_m=cfg.loss.asl.p_m, 55 | amb_k = cfg.loss.am_binary.amb_k, 56 | amb_t=cfg.loss.am_binary.amb_t, 57 | mix_precision=cfg.train.mix_precision, 58 | estimate_multilabel_thresholds=cfg.test.estimate_multilabel_thresholds) 59 | 60 | if cfg.model.type == 'classification': 61 | engine = ImageAMSoftmaxEngine( 62 | **classification_params 63 | ) 64 | elif cfg.model.type == 'multilabel': 65 | engine = MultilabelEngine( 66 | **classification_params 67 | ) 68 | elif cfg.model.type == 'multihead': 69 | engine = MultiheadEngine( 70 | **classification_params 71 | ) 72 | 73 | return engine 74 | -------------------------------------------------------------------------------- /torchreid/engine/image/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import 9 | 10 | from .am_softmax import ImageAMSoftmaxEngine 11 | from .multilabel import MultilabelEngine 12 | from .multihead import MultiheadEngine 13 | -------------------------------------------------------------------------------- /torchreid/integration/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openvinotoolkit/deep-object-reid/f695f06868de6e67acf46c4b57b66b256f3c08a1/torchreid/integration/__init__.py -------------------------------------------------------------------------------- /torchreid/integration/nncf/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 Intel Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, 10 | # software distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions 13 | # and limitations under the License. 14 | -------------------------------------------------------------------------------- /torchreid/integration/nncf/accuracy_aware_training/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2022 Intel Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, 10 | # software distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions 13 | # and limitations under the License. 14 | -------------------------------------------------------------------------------- /torchreid/integration/nncf/accuracy_aware_training/base_runner.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2022 Intel Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, 10 | # software distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions 13 | # and limitations under the License. 14 | 15 | from typing import Callable, Dict, List, Optional, Tuple, TypeVar 16 | from nncf.api.compression import CompressionAlgorithmController, CompressionStage 17 | from nncf.common.accuracy_aware_training.runner import TrainingRunner 18 | 19 | ModelType = TypeVar('ModelType') 20 | OptimizerType = TypeVar('OptimizerType') 21 | LRSchedulerType = TypeVar('LRSchedulerType') 22 | TensorboardWriterType = TypeVar('TensorboardWriterType') 23 | 24 | 25 | class BaseAccuracyAwareTrainingRunner(TrainingRunner): 26 | """ 27 | The base accuracy-aware training Runner object, 28 | initialized with the default parameters unless specified in the config. 29 | """ 30 | 31 | def __init__(self, accuracy_aware_params: Dict[str, object], verbose=True, 32 | dump_checkpoints=True): 33 | self.maximal_relative_accuracy_drop = accuracy_aware_params.get('maximal_relative_accuracy_degradation', 1.0) 34 | self.maximal_absolute_accuracy_drop = accuracy_aware_params.get('maximal_absolute_accuracy_degradation') 35 | self.maximal_total_epochs = accuracy_aware_params.get('maximal_total_epochs', 10000) 36 | self.validate_every_n_epochs = accuracy_aware_params.get('validate_every_n_epochs', 1) 37 | 38 | self.verbose = verbose 39 | self.dump_checkpoints = dump_checkpoints 40 | 41 | self.accuracy_budget = None 42 | self.is_higher_metric_better = True 43 | self._compressed_training_history = [] 44 | self._best_checkpoint = None 45 | 46 | self.training_epoch_count = 0 47 | self.cumulative_epoch_count = 0 48 | self.best_val_metric_value = 0 49 | self.loss = None 50 | 51 | def initialize_training_loop_fns(self, train_epoch_fn: Callable[[CompressionAlgorithmController, ModelType, 52 | Optional[OptimizerType], 53 | Optional[LRSchedulerType], 54 | Optional[int]], None], 55 | validate_fn: Callable[[ModelType, Optional[float]], float], 56 | configure_optimizers_fn: Callable[[], Tuple[OptimizerType, LRSchedulerType]], 57 | dump_checkpoint_fn: Callable[ 58 | [ModelType, CompressionAlgorithmController, TrainingRunner, str], None], 59 | tensorboard_writer=None, log_dir=None): 60 | self._train_epoch_fn = train_epoch_fn 61 | self._validate_fn = validate_fn 62 | self._configure_optimizers_fn = configure_optimizers_fn 63 | self._dump_checkpoint_fn = dump_checkpoint_fn 64 | self._tensorboard_writer = tensorboard_writer 65 | self._log_dir = log_dir 66 | 67 | def calculate_minimal_tolerable_accuracy(self, uncompressed_model_accuracy: float): 68 | if self.maximal_absolute_accuracy_drop is not None: 69 | self.minimal_tolerable_accuracy = uncompressed_model_accuracy - self.maximal_absolute_accuracy_drop 70 | else: 71 | self.minimal_tolerable_accuracy = uncompressed_model_accuracy * \ 72 | (1 - 0.01 * self.maximal_relative_accuracy_drop) 73 | 74 | def is_model_fully_compressed(self, compression_controller) -> bool: 75 | return compression_controller.compression_stage() == CompressionStage.FULLY_COMPRESSED 76 | 77 | def initialize_logging(self, log_dir=None, tensorboard_writer=None): 78 | pass 79 | 80 | 81 | class BaseAdaptiveCompressionLevelTrainingRunner(BaseAccuracyAwareTrainingRunner): 82 | """ 83 | The base adaptive compression level accuracy-aware training Runner object, 84 | initialized with the default parameters unless specified in the config. 85 | """ 86 | 87 | def __init__(self, accuracy_aware_params: Dict[str, object], verbose=True, 88 | minimal_compression_rate=0.05, maximal_compression_rate=0.95, 89 | dump_checkpoints=True): 90 | super().__init__(accuracy_aware_params, verbose, dump_checkpoints) 91 | 92 | self.compression_rate_step = accuracy_aware_params.get('initial_compression_rate_step', 0.1) 93 | self.step_reduction_factor = accuracy_aware_params.get('compression_rate_step_reduction_factor', 0.5) 94 | self.minimal_compression_rate_step = accuracy_aware_params.get('minimal_compression_rate_step', 0.025) 95 | self.patience_epochs = accuracy_aware_params.get('patience_epochs') 96 | self.initial_training_phase_epochs = accuracy_aware_params.get('initial_training_phase_epochs') 97 | 98 | self.minimal_compression_rate = minimal_compression_rate 99 | self.maximal_compression_rate = maximal_compression_rate 100 | 101 | self._best_checkpoints = {} 102 | self._compression_rate_target = None 103 | self.adaptive_controller = None 104 | self.was_compression_increased_on_prev_step = None 105 | 106 | @property 107 | def compression_rate_target(self): 108 | if self._compression_rate_target is None: 109 | return self.adaptive_controller.compression_rate 110 | return self._compression_rate_target 111 | 112 | @compression_rate_target.setter 113 | def compression_rate_target(self, value): 114 | self._compression_rate_target = value 115 | 116 | def get_compression_rates_with_positive_acc_budget(self) -> List[float]: 117 | return [comp_rate for (comp_rate, acc_budget) in self._compressed_training_history if acc_budget >= 0] 118 | 119 | def get_compression_rates(self) -> List[float]: 120 | return [comp_rate for (comp_rate, _) in self._compressed_training_history] 121 | -------------------------------------------------------------------------------- /torchreid/integration/nncf/compression_script_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 Intel Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, 10 | # software distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions 13 | # and limitations under the License. 14 | 15 | import logging 16 | from pprint import pformat 17 | 18 | from scripts.script_utils import build_datamanager, is_config_parameter_set_from_command_line 19 | from torchreid.integration.nncf.compression import is_nncf_state, wrap_nncf_model 20 | from torchreid.integration.nncf.config import compose_nncf_config, load_nncf_config 21 | from torchreid.integration.nncf.config import merge_dicts_and_lists_b_into_a 22 | from torchreid.utils import load_checkpoint, load_pretrained_weights 23 | 24 | logger = logging.getLogger(__name__) 25 | 26 | 27 | def get_coeff_decrease_lr_for_nncf(nncf_training_config): 28 | coef = nncf_training_config.get('coeff_decrease_lr_for_nncf') 29 | if isinstance(coef, float): 30 | return coef 31 | raise RuntimeError('The default value for coeff_decrease_lr_for_nncf is not set') 32 | 33 | 34 | def calculate_lr_for_nncf_training(cfg, initial_lr_from_checkpoint, is_initial_lr_set_from_opts): 35 | lr_from_cfg = cfg.train.lr 36 | nncf_training_config = cfg.get('nncf', {}) 37 | coeff_decrease_lr_for_nncf = get_coeff_decrease_lr_for_nncf(nncf_training_config) 38 | 39 | if is_initial_lr_set_from_opts: 40 | logger.info(f'Since initial LR is set from command line arguments, do not calculate initial LR for NNCF, ' 41 | f'taking lr from cfg, lr={lr_from_cfg}') 42 | return lr_from_cfg 43 | 44 | logger.info(f'initial_lr_from_checkpoint = {initial_lr_from_checkpoint}') 45 | if initial_lr_from_checkpoint is None: 46 | logger.info(f'The checkpoint does not contain initial lr -- will not calculate initial LR for NNCF, ' 47 | f'taking lr from cfg, lr={lr_from_cfg}') 48 | return lr_from_cfg 49 | 50 | logger.info('Try to calculate initial LR for NNCF') 51 | logger.info(f'coeff_decrease_lr_for_nncf = {coeff_decrease_lr_for_nncf}') 52 | lr = initial_lr_from_checkpoint * coeff_decrease_lr_for_nncf 53 | logger.info(f'calculated lr = {lr}') 54 | return lr 55 | 56 | 57 | def get_nncf_preset_name(enable_quantization, enable_pruning): 58 | if not isinstance(enable_quantization, bool) or not isinstance(enable_pruning, bool): 59 | return None 60 | if enable_quantization and enable_pruning: 61 | return 'nncf_quantization_pruning' 62 | if enable_quantization and not enable_pruning: 63 | return'nncf_quantization' 64 | if not enable_quantization and enable_pruning: 65 | return 'nncf_pruning' 66 | return None 67 | 68 | 69 | def make_nncf_changes_in_config(cfg, 70 | enable_quantization=False, 71 | enable_pruning=False, 72 | command_line_cfg_opts=None): 73 | 74 | if not enable_quantization and not enable_pruning: 75 | raise ValueError('None of the optimization algorithms are enabled') 76 | 77 | # default changes 78 | if cfg.lr_finder.enable: 79 | logger.info('Turn off LR finder -- it should not be used together with NNCF compression') 80 | cfg.lr_finder.enable = False 81 | 82 | cfg.nncf.enable_quantization = enable_quantization 83 | cfg.nncf.enable_pruning = enable_pruning 84 | nncf_preset = get_nncf_preset_name(enable_quantization, enable_pruning) 85 | 86 | cfg = patch_config(cfg, nncf_preset) 87 | 88 | if command_line_cfg_opts is not None: 89 | logger.info(f'applying changes to the main training config from the command line options just after that. ' 90 | f'The list of options = \n{pformat(command_line_cfg_opts)}') 91 | cfg.merge_from_list(command_line_cfg_opts) 92 | 93 | return cfg 94 | 95 | 96 | def patch_config(cfg, nncf_preset, max_acc_drop=None): 97 | # TODO: use default config here 98 | nncf_config = load_nncf_config(cfg.nncf.nncf_config_path) 99 | 100 | optimization_config = compose_nncf_config(nncf_config, [nncf_preset]) 101 | 102 | if "accuracy_aware_training" in optimization_config["nncf_config"]: 103 | # Update maximal_absolute_accuracy_degradation 104 | (optimization_config["nncf_config"]["accuracy_aware_training"] 105 | ["params"]["maximal_absolute_accuracy_degradation"]) = max_acc_drop 106 | else: 107 | logger.info("NNCF config has no accuracy_aware_training parameters") 108 | 109 | return merge_dicts_and_lists_b_into_a(cfg, optimization_config) 110 | 111 | 112 | def make_nncf_changes_in_training(model, cfg, classes, command_line_cfg_opts): 113 | if cfg.lr_finder.enable: 114 | raise RuntimeError('LR finder could not be used together with NNCF compression') 115 | 116 | if cfg.model.resume: 117 | raise NotImplementedError('Resuming NNCF training is not implemented yet') 118 | if not cfg.model.load_weights: 119 | raise RuntimeError('NNCF training should be started from a pre-trained model') 120 | checkpoint_path = cfg.model.load_weights 121 | checkpoint_dict = load_checkpoint(checkpoint_path, map_location='cpu') 122 | if is_nncf_state(checkpoint_dict): 123 | raise RuntimeError(f'The checkpoint is NNCF checkpoint at {checkpoint_path}') 124 | 125 | logger.info(f'Loading weights from {checkpoint_path}') 126 | load_pretrained_weights(model, pretrained_dict=checkpoint_dict) 127 | datamanager_for_init = build_datamanager(cfg, classes) 128 | 129 | compression_ctrl, model, nncf_metainfo = \ 130 | wrap_nncf_model(model, cfg, datamanager_for_init=datamanager_for_init) 131 | logger.info(f'Received from wrapping nncf_metainfo =\n{pformat(nncf_metainfo)}') 132 | 133 | # calculating initial LR for NNCF training 134 | lr = None 135 | initial_lr_from_checkpoint = checkpoint_dict.get('initial_lr') 136 | is_initial_lr_set_from_opts = is_config_parameter_set_from_command_line(command_line_cfg_opts, 'train.lr') 137 | lr = calculate_lr_for_nncf_training(cfg, initial_lr_from_checkpoint, 138 | is_initial_lr_set_from_opts) 139 | assert lr is not None 140 | cfg.train.lr = lr 141 | return compression_ctrl, model, cfg, lr, nncf_metainfo 142 | 143 | 144 | def make_nncf_changes_in_eval(model, cfg): 145 | _, model, _ = wrap_nncf_model(model, cfg) 146 | return model 147 | -------------------------------------------------------------------------------- /torchreid/integration/nncf/config.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 Intel Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, 10 | # software distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions 13 | # and limitations under the License. 14 | 15 | import json 16 | from copy import copy 17 | 18 | 19 | def load_nncf_config(path): 20 | assert path.endswith('.json'), ( 21 | f'Only json files are allowed as optimisation configs, provided {path}') 22 | with open(path) as f_src: 23 | nncf_config = json.load(f_src) 24 | return nncf_config 25 | 26 | 27 | def compose_nncf_config(nncf_config, enabled_options): 28 | optimisation_parts = nncf_config 29 | 30 | if 'order_of_parts' in optimisation_parts: 31 | # The result of applying the changes from optimisation parts 32 | # may depend on the order of applying the changes 33 | # (e.g. if for nncf_quantization it is sufficient to have `total_epochs=2`, 34 | # but for sparsity it is required `total_epochs=50`) 35 | # So, user can define `order_of_parts` in the optimisation_config 36 | # to specify the order of applying the parts. 37 | order_of_parts = optimisation_parts['order_of_parts'] 38 | assert isinstance(order_of_parts, list), \ 39 | 'The field "order_of_parts" in optimisation config should be a list' 40 | 41 | for part in enabled_options: 42 | assert part in order_of_parts, ( 43 | f'The part {part} is selected, but it is absent in order_of_parts={order_of_parts}') 44 | 45 | optimisation_parts_to_choose = [part for part in order_of_parts if part in enabled_options] 46 | 47 | assert 'base' in optimisation_parts, 'Error: the optimisation config does not contain the "base" part' 48 | nncf_config_part = optimisation_parts['base'] 49 | 50 | for part in optimisation_parts_to_choose: 51 | assert part in optimisation_parts, ( 52 | f'Error: the optimisation config does not contain the part "{part}"') 53 | optimisation_part_dict = optimisation_parts[part] 54 | try: 55 | nncf_config_part = merge_dicts_and_lists_b_into_a(nncf_config_part, optimisation_part_dict) 56 | except AssertionError as cur_error: 57 | err_descr = (f'Error during merging the parts of nncf configs:\n' 58 | f'the current part={part}, ' 59 | f'the order of merging parts into base is {optimisation_parts_to_choose}.\n' 60 | f'The error is:\n{cur_error}') 61 | raise RuntimeError(err_descr) from None 62 | 63 | return nncf_config_part 64 | 65 | 66 | def merge_dicts_and_lists_b_into_a(a, b): 67 | return _merge_dicts_and_lists_b_into_a(a, b, "") 68 | 69 | 70 | def _merge_dicts_and_lists_b_into_a(a, b, cur_key=None): 71 | """The function is inspired by mmcf.Config._merge_a_into_b, 72 | but it 73 | * works with usual dicts and lists and derived types 74 | * supports merging of lists (by concatenating the lists) 75 | * makes recursive merging for dict + dict case 76 | * overwrites when merging scalar into scalar 77 | Note that we merge b into a (whereas Config makes merge a into b), 78 | since otherwise the order of list merging is counter-intuitive. 79 | """ 80 | def _err_str(_a, _b, _key): 81 | if _key is None: 82 | _key_str = 'of whole structures' 83 | else: 84 | _key_str = f'during merging for key=`{_key}`' 85 | return (f'Error in merging parts of config: different types {_key_str},' 86 | f' type(a) = {type(_a)},' 87 | f' type(b) = {type(_b)}') 88 | 89 | assert isinstance(a, (dict, list)), f'Can merge only dicts and lists, whereas type(a)={type(a)}' 90 | assert isinstance(b, (dict, list)), _err_str(a, b, cur_key) 91 | assert isinstance(a, list) == isinstance(b, list), _err_str(a, b, cur_key) 92 | if isinstance(a, list): 93 | # the main diff w.r.t. mmcf.Config -- merging of lists 94 | return a + b 95 | 96 | a = copy(a) 97 | for k in b.keys(): 98 | if k not in a: 99 | a[k] = copy(b[k]) 100 | continue 101 | new_cur_key = cur_key + '.' + k if cur_key else k 102 | if isinstance(a[k], (dict, list)): 103 | a[k] = _merge_dicts_and_lists_b_into_a(a[k], b[k], new_cur_key) 104 | continue 105 | 106 | assert not isinstance(b[k], (dict, list)), _err_str(a[k], b[k], new_cur_key) 107 | 108 | # suppose here that a[k] and b[k] are scalars, just overwrite 109 | a[k] = b[k] 110 | return a 111 | -------------------------------------------------------------------------------- /torchreid/integration/nncf/engine.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 Intel Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, 10 | # software distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions 13 | # and limitations under the License. 14 | 15 | import os 16 | 17 | from torchreid.utils import ModelEmaV2 18 | from torchreid.integration.nncf.accuracy_aware_training.training_loop import ( 19 | create_accuracy_aware_training_loop 20 | ) 21 | 22 | 23 | def run_acc_aware_training_loop(engine, nncf_config, configure_optimizers_fn, stop_callback=None, perf_monitor=None): 24 | training_data = {} 25 | 26 | def dec_test(func): 27 | def inner(model, epoch): 28 | if engine.models[engine.main_model_name] != model: 29 | raise ValueError('Test function: invalid model') 30 | top1, training_data['should_save_ema_model'] = func(epoch) 31 | return top1 32 | return inner 33 | 34 | def dec_train(func): 35 | def inner(compression_controller, model, epoch, optimizer, lr_scheduler): 36 | if engine.compression_ctrl != compression_controller: 37 | raise ValueError('Train function: invalid compression controller') 38 | if engine.models[engine.main_model_name] != model: 39 | raise ValueError('Train function: invalid model') 40 | 41 | optimizers = optimizer if isinstance(optimizer, (tuple, list)) else [optimizer] 42 | lr_schedulers = lr_scheduler if isinstance(lr_scheduler, (tuple, list)) else [lr_scheduler] 43 | if len(engine.models) != len(optimizers) != len(lr_schedulers): 44 | raise ValueError('The number of optimizers and schedulers is not equal to the number of models.') 45 | 46 | for model_id, (optim, sched) in enumerate(zip(optimizers, lr_schedulers)): 47 | model_name = 'main_model' if model_id == 0 else f'aux_model_{model_id}' 48 | engine.optims[model_name] = optim 49 | engine.scheds[model_name] = sched 50 | 51 | engine.epoch = epoch 52 | return func(stop_callback=stop_callback, perf_monitor=perf_monitor) 53 | return inner 54 | 55 | def dec_update_learning_rate(func): 56 | def inner(lr_scheduler, epoch, accuracy, loss): 57 | lr_schedulers = lr_scheduler if isinstance(lr_scheduler, (tuple, list)) else [lr_scheduler] 58 | if len(engine.models) != len(lr_schedulers): 59 | raise ValueError('The number of schedulers is not equal to the number of models.') 60 | 61 | for model_id, sched in enumerate(lr_schedulers): 62 | model_name = 'main_model' if model_id == 0 else f'aux_model_{model_id}' 63 | engine.scheds[model_name] = sched 64 | 65 | engine.epoch = epoch 66 | target_metric = accuracy if engine.target_metric == 'test_acc' else loss 67 | func(output_avg_metric=target_metric) 68 | return inner 69 | 70 | def dec_early_stopping(func): 71 | def inner(accuracy): 72 | should_exit, _ = func(accuracy) 73 | return engine.early_stopping and should_exit 74 | return inner 75 | 76 | def dec_dump_checkpoint(func): 77 | def inner(model, compression_controller, _, log_dir): 78 | if engine.compression_ctrl != compression_controller: 79 | raise ValueError('Dump checkpoint function: invalid compression controller') 80 | if engine.models[engine.main_model_name] != model: 81 | raise ValueError('Dump checkpoint function: invalid model') 82 | func(0, log_dir, is_best=False, should_save_ema_model=training_data['should_save_ema_model']) 83 | return os.path.join(log_dir, 'latest.pth') 84 | return inner 85 | 86 | def reset_training(): 87 | for model_name, model in engine.models.items(): 88 | if model_name == engine.main_model_name and engine.use_ema_decay: 89 | engine.ema_model = ModelEmaV2(model, decay=engine.ema_model.decay) 90 | 91 | engine.best_metric = 0.0 92 | if hasattr(engine, 'prev_smooth_accuracy'): 93 | engine.prev_smooth_accuracy = 0.0 94 | 95 | return configure_optimizers_fn() 96 | 97 | acc_aware_training_loop = create_accuracy_aware_training_loop(nncf_config, engine.compression_ctrl, 98 | lr_updates_needed=False, dump_checkpoints=False) 99 | 100 | engine.train_data_loader = engine.train_loader 101 | engine.max_epoch = acc_aware_training_loop.runner.maximal_total_epochs 102 | 103 | validate_fn = dec_test(engine.test) 104 | train_fn = dec_train(engine.train) 105 | early_stopping_fn = dec_early_stopping(engine.exit_on_plateau_and_choose_best) 106 | dump_checkpoint_fn = dec_dump_checkpoint(engine.save_model) 107 | update_learning_rate_fn = dec_update_learning_rate(engine.update_lr) 108 | 109 | model = acc_aware_training_loop.run(engine.models[engine.main_model_name], 110 | train_epoch_fn=train_fn, 111 | validate_fn=validate_fn, 112 | configure_optimizers_fn=reset_training, 113 | early_stopping_fn=early_stopping_fn, 114 | dump_checkpoint_fn=dump_checkpoint_fn, 115 | update_learning_rate_fn=update_learning_rate_fn) 116 | return model 117 | -------------------------------------------------------------------------------- /torchreid/losses/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, division, print_function 9 | 10 | from .am_softmax import AMSoftmaxLoss, AngleSimpleLinear 11 | from .asl import AsymmetricLoss, AMBinaryLoss 12 | from .cross_entropy_loss import CrossEntropyLoss 13 | -------------------------------------------------------------------------------- /torchreid/losses/am_softmax.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2020-2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import math 18 | 19 | import torch 20 | from torch import nn 21 | import torch.nn.functional as F 22 | from torch.nn import Parameter 23 | 24 | 25 | class AngleSimpleLinear(nn.Module): 26 | """Computes cos of angles between input vectors and weights vectors""" 27 | 28 | def __init__(self, in_features, out_features): 29 | super().__init__() 30 | self.in_features = in_features 31 | self.out_features = out_features 32 | # create proxy weights 33 | self.weight = Parameter(torch.Tensor(in_features, out_features)) 34 | self.weight.data.normal_().renorm_(2, 1, 1e-5).mul_(1e5) 35 | 36 | def forward(self, x): 37 | cos_theta = F.normalize(x.view(x.shape[0], -1), dim=1).mm(F.normalize(self.weight, p=2, dim=0)) 38 | return cos_theta.clamp(-1, 1) 39 | 40 | def get_centers(self): 41 | return torch.t(self.weight) 42 | 43 | 44 | def focal_loss(input_values, gamma): 45 | """Computes the focal loss""" 46 | p = torch.exp(-input_values) 47 | loss = (1 - p) ** gamma * input_values 48 | return loss.mean() 49 | 50 | 51 | class AMSoftmaxLoss(nn.Module): 52 | margin_types = ['cos', 'arc'] 53 | 54 | def __init__(self, use_gpu=True, margin_type='cos', gamma=0.0, m=0.5, 55 | s=30, conf_penalty=0.0, label_smooth=0, aug_type='', pr_product=False, 56 | symmetric_ce=False): 57 | super().__init__() 58 | self.use_gpu = use_gpu 59 | self.conf_penalty = conf_penalty 60 | self.label_smooth = label_smooth 61 | self.aug_type = aug_type 62 | self.pr_product = pr_product 63 | 64 | assert margin_type in AMSoftmaxLoss.margin_types 65 | self.margin_type = margin_type 66 | assert gamma >= 0 67 | self.gamma = gamma 68 | assert m >= 0 69 | self.m = m 70 | assert s > 0 71 | self.scale = s 72 | self.cos_m = math.cos(m) 73 | self.sin_m = math.sin(m) 74 | self.th = math.cos(math.pi - m) 75 | self.symmetric_ce = symmetric_ce 76 | self.class_margins = self.m 77 | 78 | @staticmethod 79 | def _valid(params): 80 | if isinstance(params, (list, tuple)): 81 | for p in params: 82 | if p is None: 83 | return False 84 | else: 85 | if params is None: 86 | return False 87 | return True 88 | 89 | def forward(self, cos_theta, target, aug_index=None, lam=None, scale=None): 90 | """ 91 | Args: 92 | cos_theta (torch.Tensor): prediction matrix (before softmax) with 93 | shape (batch_size, num_classes). 94 | target (torch.LongTensor): ground truth labels with shape (batch_size). 95 | Each position contains the label index. 96 | iteration (int): current iteration 97 | """ 98 | logits_aug_avai = self._valid([self.aug_type, aug_index, lam]) # augmentations like fmix, cutmix, augmix 99 | self.scale = scale if scale else self.scale # different scale for different models 100 | if logits_aug_avai: 101 | targets1 = torch.zeros(cos_theta.size(), device=target.device).scatter_(1, target.detach().unsqueeze(1), 1) 102 | targets2 = targets1[aug_index] 103 | new_targets = targets1 * lam + targets2 * (1 - lam) 104 | # in case if target label changed 105 | target = new_targets.argmax(dim=1) 106 | 107 | if self.pr_product: 108 | pr_alpha = torch.sqrt(1.0 - cos_theta.pow(2.0)) 109 | cos_theta = pr_alpha.detach() * cos_theta + cos_theta.detach() * (1.0 - pr_alpha) 110 | 111 | one_hot_target = torch.zeros_like(cos_theta, dtype=torch.uint8) 112 | one_hot_target.scatter_(1, target.data.view(-1, 1), 1) 113 | # change margins accordingly 114 | self.class_margins *= one_hot_target 115 | 116 | if self.margin_type == 'cos': 117 | phi_theta = cos_theta - self.class_margins 118 | else: 119 | self.cos_m *= one_hot_target 120 | self.sin_m *= one_hot_target 121 | sine = torch.sqrt(1.0 - torch.pow(cos_theta, 2)) 122 | phi_theta = cos_theta * self.cos_m - sine * self.sin_m 123 | phi_theta = torch.where(cos_theta > self.th, phi_theta, cos_theta - self.sin_m * self.class_margins) 124 | 125 | output = phi_theta 126 | 127 | if self.gamma == 0.0: 128 | output *= self.scale 129 | 130 | if self.label_smooth > 0: 131 | assert not self.aug_type 132 | target = torch.zeros(output.size(), device=target.device).scatter_(1, target.detach().unsqueeze(1), 1) 133 | num_classes = output.size(1) 134 | target = (1.0 - self.label_smooth) * target + self.label_smooth / float(num_classes) 135 | losses = (- target * F.log_softmax(output, dim=1)).sum(dim=1) 136 | 137 | elif logits_aug_avai: 138 | losses = (- new_targets * F.log_softmax(output, dim=1)).sum(dim=1) 139 | 140 | else: 141 | losses = F.cross_entropy(output, target, reduction='none') 142 | 143 | if self.symmetric_ce: 144 | all_probs = F.softmax(output, dim=-1) 145 | target_probs = all_probs[torch.arange(target.size(0), device=target.device), target] 146 | losses += 4.0 * (1.0 - target_probs) 147 | 148 | if self.conf_penalty > 0.0: 149 | probs = F.softmax(output, dim=1) 150 | log_probs = F.log_softmax(output, dim=1) 151 | entropy = torch.sum(-probs * log_probs, dim=1) 152 | 153 | losses = F.relu(losses - self.conf_penalty * entropy) 154 | 155 | with torch.no_grad(): 156 | nonzero_count = max(losses.nonzero().size(0), 1) 157 | 158 | return losses.sum() / nonzero_count 159 | 160 | return focal_loss(F.cross_entropy(self.scale * output, target, reduction='none'), self.gamma) 161 | -------------------------------------------------------------------------------- /torchreid/losses/asl.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | 5 | class AsymmetricLoss(nn.Module): 6 | ''' Notice - optimized version, minimizes memory allocation and gpu uploading, 7 | favors inplace operations''' 8 | 9 | def __init__(self, gamma_neg=4, gamma_pos=0, 10 | probability_margin=0.05, eps=1e-8, 11 | label_smooth=0.): 12 | super().__init__() 13 | 14 | self.gamma_neg = gamma_neg 15 | self.gamma_pos = gamma_pos 16 | self.label_smooth = label_smooth 17 | self.clip = probability_margin 18 | self.eps = eps 19 | 20 | # prevent memory allocation and gpu uploading every iteration, and encourages inplace operations 21 | self.anti_targets = self.xs_pos = self.xs_neg = self.asymmetric_w = self.loss = None 22 | 23 | def forward(self, inputs, targets, scale=1., aug_index=None, lam=None): 24 | """" 25 | Parameters 26 | ---------- 27 | inputs: input logits 28 | targets: targets (multi-label binarized vector. Elements < 0 are ignored) 29 | """ 30 | inputs, targets = filter_input(inputs, targets) 31 | if inputs.shape[0] == 0: 32 | return 0. 33 | 34 | targets = label_smoothing(targets, self.label_smooth) 35 | self.anti_targets = 1 - targets 36 | 37 | # Calculating Probabilities 38 | self.xs_pos = torch.sigmoid(scale * inputs) 39 | self.xs_neg = 1.0 - self.xs_pos 40 | 41 | # Asymmetric Clipping 42 | if self.clip is not None and self.clip > 0: 43 | self.xs_neg.add_(self.clip).clamp_(max=1) 44 | 45 | # Basic BCE calculation 46 | self.loss = targets * torch.log(self.xs_pos.clamp(min=self.eps)) 47 | self.loss.add_(self.anti_targets * torch.log(self.xs_neg.clamp(min=self.eps))) 48 | 49 | # Asymmetric Focusing 50 | if self.gamma_neg > 0 or self.gamma_pos > 0: 51 | self.xs_pos = self.xs_pos * targets 52 | self.xs_neg = self.xs_neg * self.anti_targets 53 | self.asymmetric_w = torch.pow(1 - self.xs_pos - self.xs_neg, 54 | self.gamma_pos * targets + self.gamma_neg * self.anti_targets) 55 | self.loss *= self.asymmetric_w 56 | 57 | # sum reduction over batch 58 | return - self.loss.sum() 59 | 60 | 61 | class AMBinaryLoss(nn.Module): 62 | def __init__(self, m=0.35, k=0.8, t=1, s=30, 63 | eps=1e-8, label_smooth=0., gamma_neg=0, 64 | gamma_pos=0, probability_margin=0.05): 65 | 66 | super().__init__() 67 | self.gamma_neg = gamma_neg 68 | self.gamma_pos = gamma_pos 69 | self.asymmetric_focus = gamma_neg > 0 or gamma_pos > 0 70 | self.eps = eps 71 | self.label_smooth = label_smooth 72 | self.m = m 73 | self.k = k 74 | self.s = s 75 | self.clip = probability_margin 76 | 77 | # prevent memory allocation and gpu uploading every iteration, and encourages inplace operations 78 | self.anti_targets = self.xs_pos = self.xs_neg = self.asymmetric_w = self.loss = None 79 | 80 | def forward(self, cos_theta, targets, scale=None, aug_index=None, lam=None): 81 | """" 82 | Parameters 83 | ---------- 84 | cos_theta: dot product between normalized features and proxies 85 | targets: targets (multi-label binarized vector. Elements < 0 are ignored) 86 | """ 87 | logits_aug_avai = aug_index is not None and lam is not None # augmentations like fmix, cutmix, mixup 88 | 89 | if logits_aug_avai: 90 | aug_targets = targets[aug_index] 91 | new_targets = targets * lam + aug_targets * (1 - lam) 92 | targets = new_targets.clamp(0, 1) 93 | 94 | self.s = scale if scale else self.s 95 | 96 | cos_theta, targets = filter_input(cos_theta, targets) 97 | if cos_theta.shape[0] == 0: 98 | return 0. 99 | 100 | targets = label_smoothing(targets, self.label_smooth) 101 | self.anti_targets = 1 - targets 102 | 103 | # Calculating Probabilities 104 | xs_pos = torch.sigmoid(self.s * (cos_theta - self.m)) 105 | xs_neg = torch.sigmoid(- self.s * (cos_theta + self.m)) 106 | 107 | # SphereFace2 balancing coefficients 108 | balance_koeff_pos = self.k / self.s 109 | balance_koeff_neg = (1 - self.k) / self.s 110 | 111 | if self.asymmetric_focus: 112 | # Asymmetric Probability Shifting 113 | if self.clip is not None and self.clip > 0: 114 | xs_neg = (xs_neg + self.clip).clamp(max=1) 115 | pt0 = xs_neg * targets 116 | pt1 = xs_pos * (1 - targets) 117 | pt = pt0 + pt1 118 | one_sided_gamma = self.gamma_pos * targets + self.gamma_neg * (1 - targets) 119 | # P_pos ** gamm_neg * (...) + P_neg ** gamma_pos * (...) 120 | one_sided_w = torch.pow(pt, one_sided_gamma) 121 | 122 | self.loss = balance_koeff_pos * targets * torch.log(xs_pos) 123 | self.loss.add_(balance_koeff_neg * self.anti_targets * torch.log(xs_neg)) 124 | 125 | if self.asymmetric_focus: 126 | self.loss *= one_sided_w 127 | 128 | return - self.loss.sum() 129 | 130 | 131 | def label_smoothing(targets, smooth_degree): 132 | if smooth_degree > 0: 133 | targets = targets * (1 - smooth_degree) 134 | targets[targets == 0] = smooth_degree 135 | return targets 136 | 137 | 138 | def filter_input(inputs, targets): 139 | valid_idx = targets >= 0 140 | inputs = inputs[valid_idx] 141 | targets = targets[valid_idx] 142 | 143 | return inputs, targets 144 | -------------------------------------------------------------------------------- /torchreid/losses/cross_entropy_loss.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, division 9 | 10 | import torch 11 | from torch import nn 12 | import torch.nn.functional as F 13 | 14 | 15 | class CrossEntropyLoss(nn.Module): 16 | r"""Cross entropy loss with label smoothing regularizer. 17 | 18 | Reference: 19 | Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016. 20 | 21 | With label smoothing, the label :math:`y` for a class is computed by 22 | 23 | .. math:: 24 | \begin{equation} 25 | (1 - \eps) \times y + \frac{\eps}{K}, 26 | \end{equation} 27 | 28 | where :math:`K` denotes the number of classes and :math:`\eps` is a weight. When 29 | :math:`\eps = 0`, the loss function reduces to the normal cross entropy. 30 | 31 | Args: 32 | num_classes (int): number of classes. 33 | eps (float, optional): weight. Default is 0.1. 34 | use_gpu (bool, optional): whether to use gpu devices. Default is True. 35 | label_smooth (bool, optional): whether to apply label smoothing. Default is True. 36 | """ 37 | 38 | def __init__(self, scale=1.0, epsilon=0.1, use_gpu=True, label_smooth=0, 39 | conf_penalty=None, penalty_scale=5.0, augmentations=None): 40 | super().__init__() 41 | 42 | self.scale = scale 43 | self.label_smooth = label_smooth 44 | self.use_gpu = use_gpu 45 | self.conf_penalty = conf_penalty 46 | self.aug = augmentations 47 | self.penalty_scale = penalty_scale 48 | 49 | def forward(self, inputs, targets, aug_index=None, lam=None, scale=None): 50 | """ 51 | Args: 52 | inputs (torch.Tensor): prediction matrix (before softmax) with 53 | shape (batch_size, num_classes). 54 | targets (torch.LongTensor): ground truth labels with shape (batch_size). 55 | Each position contains the label index. 56 | """ 57 | 58 | if inputs.size(0) == 0: 59 | return torch.zeros([], dtype=inputs.dtype, device=inputs.device) 60 | 61 | scale = scale if scale is not None and scale > 0.0 else self.scale 62 | 63 | scaled_inputs = scale * inputs 64 | log_probs = F.log_softmax(scaled_inputs, dim=1) 65 | targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1) 66 | if self.use_gpu: 67 | targets = targets.cuda() 68 | 69 | num_classes = inputs.size(1) 70 | if self.label_smooth > 0: 71 | targets = (1.0 - self.label_smooth) * targets + self.label_smooth / float(num_classes) 72 | 73 | if (self.aug and aug_index is not None and lam is not None): 74 | targets2 = targets[aug_index] 75 | targets = targets * lam + targets2 * (1 - lam) 76 | 77 | sm_loss = (- targets * log_probs).sum(dim=1) 78 | 79 | if self.conf_penalty is not None and self.conf_penalty > 0.0\ 80 | and self.penalty_scale is not None and self.penalty_scale > 0.0: 81 | probs = F.softmax(scaled_inputs, dim=1) 82 | entropy = (-probs * torch.log(probs.clamp(min=1e-12))).sum(dim=1) 83 | 84 | losses = self.penalty_scale * sm_loss - self.conf_penalty * entropy 85 | losses = losses[losses > 0.0] 86 | 87 | return losses.mean() if losses.numel() > 0 else losses.sum() 88 | 89 | return sm_loss.mean() 90 | -------------------------------------------------------------------------------- /torchreid/metrics/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import 9 | 10 | from .accuracy import accuracy, accuracy_multilabel 11 | from .classification import (evaluate_classification, show_confusion_matrix, evaluate_multilabel_classification, 12 | evaluate_multihead_classification) 13 | -------------------------------------------------------------------------------- /torchreid/metrics/accuracy.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, division, print_function 9 | 10 | import torch 11 | 12 | @torch.no_grad() 13 | def accuracy(output, target, topk=(1, )): 14 | """Computes the accuracy over the k top predictions for 15 | the specified values of k. 16 | 17 | Args: 18 | output (torch.Tensor): prediction matrix with shape (batch_size, num_classes). 19 | target (torch.LongTensor): ground truth labels with shape (batch_size). 20 | topk (tuple, optional): accuracy at top-k will be computed. For example, 21 | topk=(1, 5) means accuracy at top-1 and top-5 will be computed. 22 | 23 | Returns: 24 | list: accuracy at top-k. 25 | 26 | Examples:: 27 | >>> from torchreid import metrics 28 | >>> metrics.accuracy(output, target) 29 | """ 30 | maxk = max(topk) 31 | batch_size = max(1, target.size(0)) 32 | 33 | if isinstance(output, (tuple, list)): 34 | output = output[0] 35 | 36 | _, pred = output.topk(maxk, 1, True, True) 37 | pred = pred.t() 38 | correct = pred.eq(target.view(1, -1).expand_as(pred)) 39 | 40 | res = [] 41 | for k in topk: 42 | correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) 43 | acc = correct_k.mul_(100.0 / batch_size) 44 | res.append(acc) 45 | 46 | return res 47 | 48 | 49 | @torch.no_grad() 50 | def accuracy_multilabel(output, target, threshold=0.5): 51 | batch_size = max(1, target.size(0)) 52 | 53 | if isinstance(output, (tuple, list)): 54 | output = output[0] 55 | output = torch.sigmoid(output) 56 | 57 | pred_idx = output > threshold 58 | num_correct = (pred_idx == target).sum(dim=-1) 59 | num_correct = (num_correct == target.size(1)).sum() 60 | 61 | return num_correct / batch_size 62 | -------------------------------------------------------------------------------- /torchreid/models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import 9 | 10 | from .efficient_net_pytcv import * 11 | from .inceptionv4_pytcv import * 12 | from .mobilenetv3 import * 13 | from .ptcv_wrapper import * 14 | from .timm_wrapper import * 15 | from .q2l import * 16 | from .transformer import * 17 | from .gcn import * 18 | from .ml_decoder import * 19 | 20 | __model_factory = { 21 | # image classification models 22 | 'inceptionv4': inceptionv4_pytcv, 23 | 'mobilenetv3_small': mobilenetv3_small, 24 | 'mobilenetv3_large': mobilenetv3_large, 25 | 'mobilenetv3_large_075': mobilenetv3_large_075, 26 | 'mobilenetv3_large_150': mobilenetv3_large_150, 27 | 'mobilenetv3_large_125': mobilenetv3_large_125, 28 | 'efficientnet_b0': efficientnet_b0, 29 | 'efficientnet_b1': efficientnet_b1, 30 | 'efficientnet_b2': efficientnet_b2b, 31 | 'efficientnet_b3': efficientnet_b3b, 32 | 'efficientnet_b4': efficientnet_b4b, 33 | 'efficientnet_b5': efficientnet_b5b, 34 | 'efficientnet_b6': efficientnet_b6b, 35 | 'efficientnet_b7': efficientnet_b7b, 36 | } 37 | 38 | __model_factory = {**__model_factory, **pytcv_wrapped_models, **timm_wrapped_models} 39 | 40 | 41 | def build_model(name, **kwargs): 42 | """A function wrapper for building a model. 43 | 44 | Args: 45 | name (str): model name. 46 | 47 | Returns: 48 | nn.Module 49 | 50 | Examples:: 51 | >>> from torchreid import models 52 | >>> model = models.build_model('resnet50', 751, loss='softmax') 53 | """ 54 | avai_models = list(__model_factory.keys()) 55 | if name.startswith('q2l'): 56 | backbone_name = name[4:] 57 | if backbone_name not in avai_models: 58 | raise KeyError('Unknown backbone for Q2L model: {}. Must be one of {}'.format(backbone_name, avai_models)) 59 | backbone = __model_factory[backbone_name](**kwargs) 60 | transformer = build_transformer(hidden_dim=backbone.num_features, dim_feedforward=backbone.num_features, **kwargs) 61 | model = build_q2l(backbone, transformer, hidden_dim=backbone.num_features, **kwargs) 62 | elif name.startswith('gcn'): 63 | backbone_name = name[4:] 64 | if backbone_name not in avai_models: 65 | raise KeyError('Unknown backbone for GCN model: {}. Must be one of {}'.format(backbone_name, avai_models)) 66 | backbone = __model_factory[backbone_name](**kwargs) 67 | model = build_image_gcn(backbone, **kwargs) 68 | elif name.startswith('mld'): 69 | backbone_name = name[4:] 70 | if backbone_name not in avai_models: 71 | raise KeyError('Unknown backbone for MLD model: {}. Must be one of {}'.format(backbone_name, avai_models)) 72 | backbone = __model_factory[backbone_name](**kwargs) 73 | model = build_ml_decoder_model(backbone, **kwargs) 74 | elif name not in avai_models: 75 | raise KeyError('Unknown model: {}. Must be one of {}'.format(name, avai_models)) 76 | else: 77 | model = __model_factory[name](**kwargs) 78 | return model 79 | -------------------------------------------------------------------------------- /torchreid/models/ptcv_wrapper.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2020-2021 Intel Corporation 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | http://www.apache.org/licenses/LICENSE-2.0 7 | Unless required by applicable law or agreed to in writing, software 8 | distributed under the License is distributed on an "AS IS" BASIS, 9 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 10 | See the License for the specific language governing permissions and 11 | limitations under the License. 12 | """ 13 | 14 | from functools import partial 15 | 16 | import torch.nn as nn 17 | from pytorchcv.model_provider import get_model 18 | 19 | from torchreid.losses import AngleSimpleLinear 20 | from .common import ModelInterface 21 | 22 | __all__ = ['pytcv_wrapped_models'] 23 | 24 | AVAI_MODELS = {"mobilenetv2_large" : "mobilenetv2_w1"} 25 | 26 | class PTCVModel(ModelInterface): 27 | def __init__(self, 28 | model_name, 29 | num_classes, 30 | loss='softmax', 31 | IN_first=False, 32 | pooling_type='avg', 33 | **kwargs): 34 | super().__init__(**kwargs) 35 | self.pooling_type = pooling_type 36 | self.loss = loss 37 | assert isinstance(num_classes, int) 38 | 39 | model = get_model(model_name, num_classes=1000, pretrained=self.pretrained) 40 | assert hasattr(model, 'features') and isinstance(model.features, nn.Sequential) 41 | self.features = model.features 42 | self.features = self.features[:-1] # remove pooling, since it can have a fixed size 43 | if self.loss not in ['am_softmax']: 44 | self.output_conv = nn.Conv2d(in_channels=model.output.in_channels, out_channels=num_classes, kernel_size=1, stride=1, bias=False) 45 | else: 46 | self.output_conv = AngleSimpleLinear(model.output.in_channels, num_classes) 47 | self.num_features = self.num_head_features = model.output.in_channels 48 | 49 | self.input_IN = nn.InstanceNorm2d(3, affine=True) if IN_first else None 50 | 51 | def forward(self, x, return_featuremaps=False): 52 | if self.input_IN is not None: 53 | x = self.input_IN(x) 54 | 55 | y = self.features(x) 56 | if return_featuremaps: 57 | return y 58 | 59 | glob_features = self._glob_feature_vector(y, self.pooling_type, reduce_dims=False) 60 | 61 | logits = self.output_conv(glob_features).view(x.shape[0], -1) 62 | 63 | if not self.training: 64 | return [logits] 65 | return tuple([logits]) 66 | 67 | 68 | pytcv_wrapped_models = {dor_name : partial(PTCVModel, model_name=pytcv_name) for dor_name, pytcv_name in AVAI_MODELS.items()} 69 | -------------------------------------------------------------------------------- /torchreid/models/q2l.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021 SlongLiu 2 | # SPDX-License-Identifier: MIT 3 | 4 | import torch 5 | import torch.nn as nn 6 | import math 7 | 8 | from .transformer import build_position_encoding 9 | from .common import ModelInterface 10 | from torch.cuda.amp import autocast 11 | 12 | from torchreid.losses import AngleSimpleLinear 13 | 14 | __all__ = ['build_q2l'] 15 | 16 | class GroupWiseLinear(nn.Module): 17 | def __init__(self, num_class, hidden_dim, use_bias=True): 18 | super().__init__() 19 | self.num_class = num_class 20 | self.hidden_dim = hidden_dim 21 | self.use_bias = use_bias 22 | 23 | self.weight = nn.Parameter(torch.Tensor(1, num_class, hidden_dim)) 24 | if use_bias: 25 | self.bias = nn.Parameter(torch.Tensor(1, num_class)) 26 | self.reset_parameters() 27 | 28 | def reset_parameters(self): 29 | stdv = 1. / math.sqrt(self.weight.size(2)) 30 | for i in range(self.num_class): 31 | self.weight[0][i].data.uniform_(-stdv, stdv) 32 | if self.use_bias: 33 | for i in range(self.num_class): 34 | self.bias[0][i].data.uniform_(-stdv, stdv) 35 | 36 | def forward(self, x): 37 | # x: B,K,d 38 | x = (self.weight * x).sum(-1) 39 | if self.use_bias: 40 | x = x + self.bias 41 | return x 42 | 43 | 44 | class BackboneWrapper(nn.Module): 45 | def __init__(self, backbone, position_embedding): 46 | super().__init__() 47 | self.backbone = backbone 48 | self.position_embedding = position_embedding 49 | self.num_channels = backbone.get_num_features() 50 | 51 | def forward(self, input): 52 | out = self.backbone(input, return_featuremaps=True) 53 | pos = self.position_embedding(out).to(out.dtype) 54 | return [out], [pos] 55 | 56 | 57 | class Query2Label(ModelInterface): 58 | def __init__(self, 59 | backbone, 60 | transfomer, 61 | num_classes=80, 62 | pretrain=False, 63 | **kwargs): 64 | """[summary] 65 | 66 | Args: 67 | backbone ([type]): backbone model. 68 | transfomer ([type]): transformer model. 69 | num_classes ([type]): number of classes. (80 for MSCOCO). 70 | """ 71 | super().__init__(**kwargs) 72 | self.backbone = backbone 73 | self.transformer = transfomer 74 | self.num_class = num_classes 75 | assert self.loss in ['asl', 'bce', 'am_binary'], "Q2L supports only ASL, BCE pr AM Binary losses" 76 | 77 | hidden_dim = transfomer.get_hidden_dim() 78 | backbone_features = backbone.num_channels 79 | self.input_proj = nn.Conv2d(backbone_features, hidden_dim, kernel_size=1) 80 | self.query_embed = nn.Embedding(num_classes, hidden_dim) 81 | if self.loss == "am_binary": 82 | self.fc = AngleSimpleLinear(hidden_dim, num_classes) 83 | else: 84 | assert self.loss in ['asl', 'bce'] 85 | self.fc = GroupWiseLinear(num_classes, hidden_dim, use_bias=True) 86 | 87 | def forward(self, input): 88 | with autocast(enabled=self.mix_precision): 89 | src, pos = self.backbone(input) 90 | src, pos = src[-1], pos[-1] 91 | 92 | query_input = self.query_embed.weight 93 | hs = self.transformer(self.input_proj(src), query_input, pos)[0] # B,K,d 94 | logits = self.fc(hs[-1]) 95 | if self.similarity_adjustment: 96 | logits = self.sym_adjust(logits, self.amb_t) 97 | 98 | if not self.training: 99 | return [logits] 100 | 101 | elif self.loss in ['asl', 'bce', 'am_binary']: 102 | out_data = [logits] 103 | else: 104 | raise KeyError("Unsupported loss: {}".format(self.loss)) 105 | 106 | return tuple(out_data) 107 | 108 | def get_config_optim(self, lrs): 109 | parameters = [ 110 | {'params': self.backbone.named_parameters()}, 111 | {'params': self.fc.named_parameters()}, 112 | {'params': self.input_proj.parameters(), 'weight_decay': 0.}, 113 | {'params': self.query_embed.parameters(), 'weight_decay': 0.} 114 | ] 115 | if isinstance(lrs, list): 116 | assert len(lrs) == len(parameters) 117 | for lr, param_dict in zip(lrs, parameters): 118 | param_dict['lr'] = lr 119 | else: 120 | assert isinstance(lrs, float) 121 | for param_dict in parameters: 122 | param_dict['lr'] = lrs 123 | 124 | return parameters 125 | 126 | 127 | def build_q2l(backbone, transformer, hidden_dim=2048, pretrain=False, input_size=448, **kwargs): 128 | position_emb = build_position_encoding(hidden_dim=hidden_dim, img_size=input_size) 129 | wrapped_model = BackboneWrapper(backbone, position_emb) 130 | model = Query2Label( 131 | backbone=wrapped_model, 132 | transfomer=transformer, 133 | pretrain=pretrain, 134 | **kwargs 135 | ) 136 | 137 | return model 138 | -------------------------------------------------------------------------------- /torchreid/models/timm_wrapper.py: -------------------------------------------------------------------------------- 1 | import timm 2 | 3 | from torchreid.losses import AngleSimpleLinear 4 | from .common import ModelInterface 5 | import torchreid.utils as utils 6 | from torchreid.ops import Dropout 7 | from torch import nn 8 | from torch.cuda.amp import autocast 9 | 10 | 11 | __all__ = ["timm_wrapped_models", "TimmModelsWrapper"] 12 | AVAI_MODELS = { 13 | 'mobilenetv3_large_21k' : 'mobilenetv3_large_100_miil_in21k', 14 | 'mobilenetv3_large_1k' : 'mobilenetv3_large_100_miil', 15 | 'tresnet_m' : 'tresnet_m', 16 | 'tresnet_l' : 'tresnet_l', 17 | 'efficientnetv2_s_21k': 'tf_efficientnetv2_s_in21k', 18 | 'efficientnetv2_s_1k': 'tf_efficientnetv2_s_in21ft1k', 19 | 'efficientnetv2_m_21k': 'tf_efficientnetv2_m_in21k', 20 | 'efficientnetv2_m_1k': 'tf_efficientnetv2_m_in21ft1k', 21 | 'efficientnetv2_l_21k': 'tf_efficientnetv2_l_in21k', 22 | 'efficientnetv2_l_1k': 'tf_efficientnetv2_l_in21ft1k', 23 | 'efficientnetv2_b0' : 'tf_efficientnetv2_b0', 24 | 'resnet101' : "tv_resnet101" 25 | } 26 | 27 | class TimmModelsWrapper(ModelInterface): 28 | def __init__(self, 29 | model_name, 30 | pretrained=False, 31 | dropout_cls = None, 32 | pooling_type='avg', 33 | emb_dim=1680, 34 | **kwargs): 35 | super().__init__(**kwargs) 36 | self.pretrained = pretrained 37 | self.is_mobilenet = True if model_name in ["mobilenetv3_large_100_miil_in21k", "mobilenetv3_large_100_miil"] else False 38 | self.model = timm.create_model(model_name, 39 | pretrained=pretrained, 40 | num_classes=self.num_classes) 41 | self.num_head_features = self.model.num_features 42 | self.num_features = (self.model.conv_head.in_channels if self.is_mobilenet 43 | else self.model.num_features) 44 | self.dropout = Dropout(**dropout_cls) 45 | self.pooling_type = pooling_type 46 | if "am_softmax" in self.loss or "am_binary" in self.loss: 47 | self.model.act2 = nn.PReLU() 48 | self.model.classifier = AngleSimpleLinear(self.num_head_features, self.num_classes) 49 | else: 50 | assert "softmax" in self.loss or "asl" in self.loss 51 | self.model.classifier = self.model.get_classifier() 52 | 53 | def forward(self, x, return_featuremaps=False, return_all=False, apply_scale=False, **kwargs): 54 | with autocast(enabled=self.mix_precision): 55 | y = self.extract_features(x) 56 | if return_featuremaps: 57 | return y 58 | glob_features = self._glob_feature_vector(y, self.pooling_type, reduce_dims=False) 59 | logits = self.infer_head(glob_features) 60 | if self.similarity_adjustment: 61 | logits = self.sym_adjust(logits, self.amb_t) 62 | if apply_scale: 63 | logits *= self.scale 64 | 65 | if return_all: 66 | saliency_map = self._extract_saliency_map(y) 67 | return [(logits, saliency_map, glob_features)] 68 | if not self.training: 69 | return [logits] 70 | return tuple([logits]) 71 | 72 | def extract_features(self, x): 73 | return self.model.forward_features(x) 74 | 75 | def infer_head(self, x): 76 | if self.is_mobilenet: 77 | x = self.model.act2(self.model.conv_head(x)) 78 | self.dropout(x) 79 | return self.model.classifier(x.view(x.shape[0], -1)) 80 | 81 | def get_config_optim(self, lrs): 82 | parameters = [ 83 | {'params': self.model.named_parameters()}, 84 | ] 85 | if isinstance(lrs, list): 86 | assert len(lrs) == len(parameters) 87 | for lr, param_dict in zip(lrs, parameters): 88 | param_dict['lr'] = lr 89 | else: 90 | assert isinstance(lrs, float) 91 | for param_dict in parameters: 92 | param_dict['lr'] = lrs 93 | 94 | return parameters 95 | 96 | class ModelFactory: 97 | def __init__(self, model_name) -> None: 98 | self.model_name = model_name 99 | 100 | def __call__(self, **kwargs): 101 | """ 102 | Constructs a timm model 103 | """ 104 | net = TimmModelsWrapper(self.model_name, **kwargs) 105 | return net 106 | 107 | timm_wrapped_models = {dor_name : ModelFactory(model_name=timm_name) for dor_name, timm_name in AVAI_MODELS.items()} 108 | -------------------------------------------------------------------------------- /torchreid/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2020-2021 Intel Corporation 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | 5 | from .data_parallel import DataParallel 6 | from .dropout import Dropout 7 | from .self_challenging import RSC, rsc 8 | -------------------------------------------------------------------------------- /torchreid/ops/data_parallel.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2020-2021 Intel Corporation 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | 5 | from itertools import chain 6 | 7 | import torch 8 | from torch import nn 9 | 10 | 11 | def map_device(outputs, output_device): 12 | if isinstance(outputs, torch.Tensor): 13 | return outputs.to(output_device) 14 | if isinstance(outputs, (tuple, list)): 15 | return [map_device(o, output_device) for o in outputs] 16 | if isinstance(outputs, dict): 17 | return {k: map_device(v, output_device) for k, v in outputs.items()} 18 | 19 | raise ValueError(f'Unknown output type: {type(outputs)}') 20 | 21 | 22 | class DataParallel(nn.DataParallel): 23 | def forward(self, *inputs, **kwargs): 24 | if not self.device_ids: 25 | return self.module(*inputs, **kwargs) 26 | 27 | for t in chain(self.module.parameters(), self.module.buffers()): 28 | if t.device != self.src_device_obj: 29 | raise RuntimeError("module must have its parameters and buffers " 30 | f"on device {self.src_device_obj} (device_ids[0]) but found one of " 31 | f"them on device: {t.device}") 32 | 33 | inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids) 34 | 35 | if len(self.device_ids) == 1: 36 | outputs = self.module(*inputs[0], **kwargs[0]) 37 | return map_device(outputs, self.output_device) 38 | 39 | replicas = self.replicate(self.module, self.device_ids[:len(inputs)]) 40 | outputs = self.parallel_apply(replicas, inputs, kwargs) 41 | 42 | return self.gather(outputs, self.output_device) 43 | -------------------------------------------------------------------------------- /torchreid/ops/dropout.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2020-2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import torch 18 | from torch import nn 19 | import torch.nn.functional as F 20 | 21 | 22 | class Dropout(nn.Module): 23 | DISTRIBUTIONS = ['none', 'bernoulli', 'gaussian', 'infodrop'] 24 | 25 | def __init__(self, p=0.1, mu=0.1, sigma=0.03, dist='bernoulli', kernel=3, temperature=0.2): 26 | super().__init__() 27 | 28 | self.dist = dist 29 | assert self.dist in Dropout.DISTRIBUTIONS 30 | 31 | self.p = float(p) 32 | assert 0. <= self.p <= 1. 33 | 34 | self.mu = float(mu) 35 | self.sigma = float(sigma) 36 | assert self.sigma > 0. 37 | 38 | self.kernel = kernel 39 | assert self.kernel >= 3 40 | self.temperature = temperature 41 | assert self.temperature > 0.0 42 | 43 | def forward(self, x, x_original=None): 44 | if not self.training: 45 | return x 46 | 47 | if self.dist == 'bernoulli': 48 | out = F.dropout(x, self.p, self.training) 49 | elif self.dist == 'gaussian': 50 | with torch.no_grad(): 51 | soft_mask = x.new_empty(x.size()).normal_(self.mu, self.sigma).clamp_(0., 1.) 52 | 53 | scale = 1. / self.mu 54 | out = scale * soft_mask * x 55 | elif self.dist == 'infodrop': 56 | assert x_original is not None 57 | 58 | out = info_dropout(x_original, self.kernel, x, self.p, self.temperature) 59 | else: 60 | out = x 61 | 62 | return out 63 | 64 | 65 | def info_dropout(in_features, kernel, out_features, drop_rate, temperature=0.2, eps=1e-12): 66 | assert isinstance(kernel, int) 67 | assert kernel % 2 == 1 68 | 69 | in_shape = in_features.size() 70 | assert len(in_shape) in (4, 5) 71 | 72 | with torch.no_grad(): 73 | b, c, h, w = in_shape 74 | out_mask_shape = b, 1, h, w 75 | in_features = in_features.reshape(-1, c, h, w) 76 | 77 | padding = (kernel - 1) // 2 78 | unfolded_features = F.unfold(in_features, kernel, padding=padding) 79 | unfolded_features = unfolded_features.view(b, c, kernel * kernel, -1) 80 | 81 | distances = ((unfolded_features - in_features.view(-1, c, 1, h * w)) ** 2).sum(dim=1) 82 | weights = (0.5 * distances / distances.mean(dim=(1, 2), keepdim=True).clamp_min(eps)).neg().exp() 83 | 84 | middle = kernel * kernel // 2 85 | log_info = (weights[:, :middle].sum(dim=1) + weights[:, (middle + 1):].sum(dim=1) + eps).log() 86 | 87 | prob_weights = (1. / float(temperature) * log_info).exp() + eps 88 | probs = prob_weights / prob_weights.sum(dim=-1, keepdim=True) 89 | 90 | drop_num_samples = max(1, int(drop_rate * float(probs.size(-1)))) 91 | drop_indices = torch.multinomial(probs, num_samples=drop_num_samples, replacement=True) 92 | 93 | out_mask = torch.ones_like(probs) 94 | out_mask[torch.arange(out_mask.size(0), device=out_mask.device).view(-1, 1), drop_indices] = 0.0 95 | 96 | out_scale = 1.0 / (1.0 - drop_rate) 97 | out = out_scale * out_features * out_mask.view(out_mask_shape) 98 | 99 | return out 100 | -------------------------------------------------------------------------------- /torchreid/ops/self_challenging.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2020-2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | ''' 16 | 17 | import numpy as np 18 | import torch 19 | from torch import nn 20 | from torch.autograd import grad 21 | 22 | 23 | class RSC(nn.Module): 24 | def __init__(self, retain_p): 25 | super().__init__() 26 | 27 | self.retain_p = float(retain_p) 28 | assert 0. < self.retain_p < 1. 29 | 30 | def forward(self, features, scores, labels): 31 | return rsc(features, scores, labels, self.retain_p) 32 | 33 | 34 | def rsc(features, scores, labels, retain_p=0.67, retain_batch=0.67): 35 | """Representation Self-Challenging module (RSC). 36 | Based on the paper: https://arxiv.org/abs/2007.02454 37 | """ 38 | 39 | batch_range = torch.arange(scores.size(0), device=scores.device) 40 | gt_scores = scores[batch_range, labels.view(-1)] 41 | z_grads = grad(outputs=gt_scores, 42 | inputs=features, 43 | grad_outputs=torch.ones_like(gt_scores), 44 | create_graph=True)[0] 45 | 46 | with torch.no_grad(): 47 | z_grads_cpu = z_grads.cpu().numpy() 48 | non_batch_axis = tuple(range(1, len(z_grads_cpu.shape))) 49 | z_grad_thresholds_cpu = np.quantile(z_grads_cpu, retain_p, axis=non_batch_axis, keepdims=True) 50 | zero_mask = z_grads > torch.from_numpy(z_grad_thresholds_cpu).to(z_grads.device) 51 | 52 | unchanged_mask = torch.randint(2, [z_grads.size(0)], dtype=torch.bool, device=z_grads.device) 53 | unchanged_mask = unchanged_mask.view((-1,) + (1,) * len(non_batch_axis)) 54 | 55 | scale = 1.0 / float(retain_p) 56 | filtered_features = scale * torch.where(zero_mask, torch.zeros_like(features), features) 57 | out_features = torch.where(unchanged_mask, features, filtered_features) 58 | # filter batch 59 | random_uniform = torch.rand(size=(features.size(0), 1, 1, 1), device=features.device) 60 | random_mask = random_uniform >= retain_batch 61 | out_features = torch.where(random_mask, out_features, features) 62 | 63 | return out_features 64 | -------------------------------------------------------------------------------- /torchreid/optim/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from .lr_scheduler import build_lr_scheduler, WarmupScheduler, ReduceLROnPlateauV2, CosineAnnealingCycleRestart 4 | from .optimizer import build_optimizer 5 | from .sam import SAM 6 | from .lr_finder import LrFinder 7 | -------------------------------------------------------------------------------- /torchreid/optim/lr_finder.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import numpy as np 4 | from torch_lr_finder import LRFinder 5 | import optuna 6 | from optuna.trial import TrialState 7 | from optuna.samplers import TPESampler, GridSampler 8 | from addict import Dict 9 | import time 10 | import datetime 11 | 12 | from torchreid.utils import get_model_attr 13 | from torchreid.optim.sam import SAM 14 | 15 | 16 | class WrappedLRFinder(LRFinder): 17 | def _train_batch(self, train_iter, accumulation_steps, non_blocking_transfer=True): 18 | self.model.train() 19 | total_loss = None # for late initialization 20 | 21 | self.optimizer.zero_grad() 22 | for _ in range(accumulation_steps): 23 | inputs, labels = next(train_iter) 24 | inputs, labels = self._move_to_device( 25 | inputs, labels, non_blocking=non_blocking_transfer 26 | ) 27 | # Forward pass 28 | loss = self.forward_pass(inputs, labels, accumulation_steps) 29 | 30 | if not isinstance(self.optimizer, SAM): 31 | self.optimizer.step() 32 | else: 33 | self.optimizer.first_step() 34 | loss = self.forward_pass(inputs, labels, accumulation_steps) 35 | self.optimizer.second_step() 36 | 37 | if total_loss is None: 38 | total_loss = loss 39 | else: 40 | total_loss += loss 41 | 42 | return total_loss.item() 43 | 44 | def forward_pass(self, inputs, targets, accumulation_steps): 45 | # Forward pass 46 | model_output = self.model(inputs) 47 | all_logits = model_output[0] if isinstance(model_output, (tuple, list)) else model_output 48 | loss = self.criterion(all_logits, targets) 49 | 50 | # Loss should be averaged in each step 51 | loss /= accumulation_steps 52 | 53 | # Backward pass 54 | loss.backward() 55 | return loss 56 | 57 | 58 | class LrFinder: 59 | def __init__(self, 60 | engine, 61 | mode='fast_ai', 62 | epochs_warmup=2, 63 | max_lr=0.03, 64 | min_lr=4e-3, 65 | step=0.001, 66 | num_epochs=3, 67 | path_to_savefig='', 68 | seed = 5, 69 | stop_callback=None, 70 | smooth_f=0.01, 71 | n_trials=30, 72 | **kwargs) -> None: 73 | r"""A pipeline for learning rate search. 74 | Args: 75 | mode (str, optional): mode for learning rate finder, "fast_ai", "grid_search", "TPE". 76 | Default is "fast_ai". 77 | max_lr (float): upper bound for leaning rate 78 | min_lr (float): lower bound for leaning rate 79 | step (float, optional): number of step for learning rate searching space. Default is 1e-3 80 | num_epochs (int, optional): number of epochs to train for each learning rate. Default is 3 81 | pretrained (bool): whether or not the model is pretrained 82 | path_to_savefig (str): if path given save plot loss/lr (only for fast_ai mode). Default: '' 83 | """ 84 | self.engine = engine 85 | main_model_name = engine.get_model_names(None)[0] 86 | self.model = engine.models[main_model_name] 87 | self.optimizer = engine.optims[main_model_name] 88 | self.model_device = next(self.model.parameters()).device 89 | self.mode = mode 90 | self.min_lr = min_lr 91 | self.max_lr = max_lr 92 | self.step = step 93 | self.n_trials = n_trials 94 | self.num_epochs = num_epochs 95 | self.path_to_savefig = path_to_savefig 96 | self.seed = seed 97 | self.stop_callback = stop_callback 98 | self.epochs_warmup = epochs_warmup 99 | self.enable_sam = engine.enable_sam 100 | self.smooth_f = smooth_f 101 | self.engine_cfg = Dict(min_lr=min_lr, max_lr=max_lr, mode=mode, step=step) 102 | search_space = np.arange(min_lr, max_lr, step) 103 | self.samplers = {'grid_search': GridSampler(search_space={'lr': search_space}), 104 | 'TPE': TPESampler(n_startup_trials=5, seed=True)} 105 | 106 | def process(self): 107 | print(f'=> Start learning rate search. Mode: {self.mode}') 108 | if self.mode == 'fast_ai': 109 | lr = self.fast_ai() 110 | return lr 111 | 112 | assert self.mode in ['grid_search', 'TPE'] 113 | lr = self.optuna_optim() 114 | return lr 115 | 116 | def fast_ai(self): 117 | criterion = self.engine.main_losses[0] 118 | 119 | if self.epochs_warmup != 0: 120 | get_model_attr(self.model, 'to')(self.model_device) 121 | print(f"Warmup the model's weights for {self.epochs_warmup} epochs") 122 | self.engine.run(max_epoch=self.epochs_warmup, lr_finder=self.engine_cfg, 123 | stop_callback=self.stop_callback, eval_freq=1) 124 | print("Finished warmuping the model. Continue to find learning rate:") 125 | 126 | # run lr finder 127 | num_iter = len(self.engine.train_loader) 128 | lr_finder = WrappedLRFinder(self.model, self.optimizer, criterion, device=self.model_device) 129 | lr_finder.range_test(self.engine.train_loader, start_lr=self.min_lr, end_lr=self.max_lr, 130 | smooth_f=self.smooth_f, num_iter=num_iter, step_mode='exp') 131 | ax, optim_lr = lr_finder.plot(suggest_lr=True) 132 | # save plot if needed 133 | if self.path_to_savefig: 134 | fig = ax.get_figure() 135 | fig.savefig(self.path_to_savefig) 136 | 137 | # reset weights and optimizer state 138 | if self.epochs_warmup != 0: 139 | self.engine.restore_model() 140 | else: 141 | lr_finder.reset() 142 | 143 | return optim_lr 144 | 145 | def optuna_optim(self): 146 | def finish_process(study): 147 | pruned_trials = study.get_trials(deepcopy=False, states=[TrialState.PRUNED]) 148 | complete_trials = study.get_trials(deepcopy=False, states=[TrialState.COMPLETE]) 149 | 150 | print("Study statistics: ") 151 | print(" Number of finished trials: ", len(study.trials)) 152 | print(" Number of pruned trials: ", len(pruned_trials)) 153 | print(" Number of complete trials: ", len(complete_trials)) 154 | 155 | print("Best trial:") 156 | trial = study.best_trial 157 | 158 | print(" Value: ", trial.value) 159 | 160 | print(" Params: ") 161 | for key, value in trial.params.items(): 162 | print(f" {key}: {value}") 163 | 164 | return trial.params['lr'] 165 | 166 | study = optuna.create_study(study_name='classification task', 167 | direction="maximize", sampler=self.samplers[self.mode]) 168 | objective_partial = partial(self.engine.run, max_epoch=self.num_epochs, 169 | lr_finder=self.engine_cfg, start_eval=0, eval_freq=1, 170 | stop_callback=self.stop_callback) 171 | objective = lambda trial: objective_partial(trial)[0] 172 | try: 173 | start_time = time.time() 174 | study.optimize(objective, n_trials=self.n_trials, timeout=None) 175 | elapsed = round(time.time() - start_time) 176 | print(f"--- learning rate estimation finished with elapsed time: {datetime.timedelta(seconds=elapsed)} ---") 177 | 178 | except KeyboardInterrupt: 179 | return finish_process(study) 180 | 181 | else: 182 | return finish_process(study) 183 | -------------------------------------------------------------------------------- /torchreid/optim/sam.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Research 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | ''' 9 | Imported from: https://github.com/google-research/sam 10 | ''' 11 | 12 | import torch 13 | 14 | 15 | class SAM(torch.optim.Optimizer): 16 | def __init__(self, params, base_optimizer, rho=0.05, adaptive=True, **kwargs): 17 | assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}" 18 | self.base_optimizer = base_optimizer 19 | defaults = dict(rho=rho, adaptive=adaptive, **self.base_optimizer.defaults) 20 | 21 | super().__init__(params, defaults) 22 | self.rho = rho 23 | self.adaptive = adaptive 24 | self.param_groups = self.base_optimizer.param_groups 25 | 26 | @torch.no_grad() 27 | def first_step(self, zero_grad=False): 28 | if self._has_overflow(self.param_groups): 29 | if zero_grad: self.zero_grad() 30 | return True 31 | 32 | grad_norm = self._grad_norm() 33 | for group in self.param_groups: 34 | scale = self.rho / (grad_norm + 1e-12) 35 | 36 | for p in group["params"]: 37 | if p.grad is None: continue 38 | self.state[p]["old_p"] = p.data.clone() 39 | e_w = (torch.pow(p, 2) if self.adaptive else 1.0) * p.grad * scale.to(p) 40 | p.add_(e_w) # climb to the local maximum "w + e(w)" 41 | 42 | if zero_grad: self.zero_grad() 43 | return False 44 | 45 | @torch.no_grad() 46 | def second_step(self, zero_grad=False): 47 | if self._has_overflow(self.param_groups): 48 | if zero_grad: self.zero_grad() 49 | return 50 | 51 | for group in self.param_groups: 52 | for p in group["params"]: 53 | if p.grad is None: continue 54 | p.data = self.state[p]["old_p"] # get back to "w" from "w + e(w)" 55 | 56 | self.base_optimizer.step() # do the actual "sharpness-aware" update 57 | 58 | if zero_grad: self.zero_grad() 59 | 60 | @torch.no_grad() 61 | def step(self): 62 | raise NotImplementedError("SAM doesn't work like the other optimizers," 63 | " you should first call `first_step` and the `second_step`;") 64 | 65 | def _grad_norm(self): 66 | shared_device = self.param_groups[0]["params"][0].device # put everything on the same device, in case of model parallelism 67 | norm = torch.norm( 68 | torch.stack([ 69 | ((torch.abs(p) if self.adaptive else 1.0) * p.grad).norm(p=2).to(shared_device) 70 | for group in self.param_groups for p in group["params"] 71 | if p.grad is not None 72 | ]), 73 | p=2 74 | ) 75 | return norm 76 | 77 | @staticmethod 78 | def _has_overflow(params): 79 | ''' Check whether the gradient overflow occurred in model parameters ''' 80 | def _has_inf_or_nan(x): 81 | try: 82 | # if x is half, the .float() incurs an additional deep copy, but it's necessary if 83 | # Pytorch's .sum() creates a one-element tensor of the same type as x 84 | # (which is true for some recent version of pytorch). 85 | cpu_sum = float(x.float().sum()) 86 | # More efficient version that can be used if .sum() returns a Python scalar 87 | # cpu_sum = float(x.sum()) 88 | except RuntimeError as instance: 89 | # We want to check if inst is actually an overflow exception. 90 | # RuntimeError could come from a different error. 91 | # If so, we still want the exception to propagate. 92 | if "value cannot be converted" not in instance.args[0]: 93 | raise 94 | return True 95 | else: 96 | if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: 97 | return True 98 | return False 99 | 100 | for group in params: 101 | for p in group["params"]: 102 | if p.grad is not None and _has_inf_or_nan(p.grad.data): 103 | return True 104 | 105 | return False 106 | -------------------------------------------------------------------------------- /torchreid/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from .tools import * 4 | from .loggers import * 5 | from .avgmeter import * 6 | from .torchtools import * 7 | from .fmix import sample_mask 8 | -------------------------------------------------------------------------------- /torchreid/utils/avgmeter.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | from __future__ import absolute_import, division 9 | from collections import defaultdict 10 | 11 | import torch 12 | 13 | __all__ = ['AverageMeter', 'MetricMeter'] 14 | 15 | 16 | class AverageMeter: 17 | """Computes and stores the average and current value. 18 | 19 | Examples:: 20 | >>> # Initialize a meter to record loss 21 | >>> losses = AverageMeter() 22 | >>> # Update meter after every minibatch update 23 | >>> losses.update(loss_value, batch_size) 24 | """ 25 | 26 | def __init__(self, enable_zeros=False): 27 | self.enable_zeros = enable_zeros 28 | 29 | self.reset() 30 | 31 | def reset(self): 32 | self.val = 0.0 33 | self.avg = 0.0 34 | self.sum = 0.0 35 | self.count = 0 36 | 37 | def update(self, val, n=1): 38 | self.val = val 39 | if n > 0: 40 | if self.enable_zeros: 41 | self._update(val, n) 42 | elif val > 0.0: 43 | self._update(val, n) 44 | 45 | def _update(self, val, n): 46 | self.sum += val * n 47 | self.count += n 48 | self.avg = self.sum / self.count 49 | 50 | 51 | class MetricMeter: 52 | """A collection of metrics. 53 | 54 | Source: https://github.com/KaiyangZhou/Dassl.pytorch 55 | 56 | Examples:: 57 | >>> # 1. Create an instance of MetricMeter 58 | >>> metric = MetricMeter() 59 | >>> # 2. Update using a dictionary as input 60 | >>> input_dict = {'loss_1': value_1, 'loss_2': value_2} 61 | >>> metric.update(input_dict) 62 | >>> # 3. Convert to string and print 63 | >>> print(str(metric)) 64 | """ 65 | 66 | def __init__(self, delimiter='\t'): 67 | self.meters = defaultdict(AverageMeter) 68 | self.delimiter = delimiter 69 | 70 | def update(self, input_dict): 71 | if input_dict is None: 72 | return 73 | 74 | if not isinstance(input_dict, dict): 75 | raise TypeError( 76 | 'Input to MetricMeter.update() must be a dictionary' 77 | ) 78 | 79 | for k, v in input_dict.items(): 80 | if isinstance(v, torch.Tensor): 81 | v = v.item() 82 | self.meters[k].update(v) 83 | 84 | def __str__(self): 85 | output_str = [] 86 | for name, meter in self.meters.items(): 87 | output_str.append( 88 | f'{name} {meter.val:.4f} ({meter.avg:.4f})' 89 | ) 90 | return self.delimiter.join(output_str) 91 | -------------------------------------------------------------------------------- /torchreid/utils/fmix.py: -------------------------------------------------------------------------------- 1 | ''' 2 | imported from https://github.com/ecs-vlc/FMix 3 | 4 | MIT License 5 | 6 | Copyright (c) 2020 Vision, Learning and Control Research Group 7 | 8 | Permission is hereby granted, free of charge, to any person obtaining a copy 9 | of this software and associated documentation files (the "Software"), to deal 10 | in the Software without restriction, including without limitation the rights 11 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | copies of the Software, and to permit persons to whom the Software is 13 | furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in all 16 | copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 24 | SOFTWARE. 25 | 26 | Copyright (C) 2020-2021 Intel Corporation 27 | SPDX-License-Identifier: Apache-2.0 28 | ''' 29 | 30 | import math 31 | import random 32 | 33 | import numpy as np 34 | from scipy.stats import beta 35 | 36 | 37 | def fftfreqnd(h, w=None, z=None): 38 | """ Get bin values for discrete fourier transform of size (h, w, z) 39 | :param h: Required, first dimension size 40 | :param w: Optional, second dimension size 41 | :param z: Optional, third dimension size 42 | """ 43 | fz = fx = 0 44 | fy = np.fft.fftfreq(h) 45 | 46 | if w is not None: 47 | fy = np.expand_dims(fy, -1) 48 | 49 | if w % 2 == 1: 50 | fx = np.fft.fftfreq(w)[: w // 2 + 2] 51 | else: 52 | fx = np.fft.fftfreq(w)[: w // 2 + 1] 53 | 54 | if z is not None: 55 | fy = np.expand_dims(fy, -1) 56 | if z % 2 == 1: 57 | fz = np.fft.fftfreq(z)[:, None] 58 | else: 59 | fz = np.fft.fftfreq(z)[:, None] 60 | 61 | return np.sqrt(fx * fx + fy * fy + fz * fz) 62 | 63 | 64 | def get_spectrum(freqs, decay_power, ch, h, w=0, z=0): 65 | """ Samples a fourier image with given size and frequencies decayed by decay power 66 | :param freqs: Bin values for the discrete fourier transform 67 | :param decay_power: Decay power for frequency decay prop 1/f**d 68 | :param ch: Number of channels for the resulting mask 69 | :param h: Required, first dimension size 70 | :param w: Optional, second dimension size 71 | :param z: Optional, third dimension size 72 | """ 73 | scale = np.ones(1) / (np.maximum(freqs, np.array([1. / max(w, h, z)])) ** decay_power) 74 | 75 | param_size = [ch] + list(freqs.shape) + [2] 76 | param = np.random.randn(*param_size) 77 | 78 | scale = np.expand_dims(scale, -1)[None, :] 79 | 80 | return scale * param 81 | 82 | 83 | def make_low_freq_image(decay, shape, ch=1): 84 | """ Sample a low frequency image from fourier space 85 | :param decay_power: Decay power for frequency decay prop 1/f**d 86 | :param shape: Shape of desired mask, list up to 3 dims 87 | :param ch: Number of channels for desired mask 88 | """ 89 | freqs = fftfreqnd(*shape) 90 | spectrum = get_spectrum(freqs, decay, ch, *shape) 91 | spectrum = spectrum[:, 0] + 1j * spectrum[:, 1] 92 | mask = np.real(np.fft.irfftn(spectrum, shape)) 93 | 94 | if len(shape) == 1: 95 | mask = mask[:1, :shape[0]] 96 | if len(shape) == 2: 97 | mask = mask[:1, :shape[0], :shape[1]] 98 | if len(shape) == 3: 99 | mask = mask[:1, :shape[0], :shape[1], :shape[2]] 100 | 101 | mask = (mask - mask.min()) 102 | mask = mask / mask.max() 103 | return mask 104 | 105 | 106 | def sample_lam(alpha, reformulate=False): 107 | """ Sample a lambda from symmetric beta distribution with given alpha 108 | :param alpha: Alpha value for beta distribution 109 | :param reformulate: If True, uses the reformulation of [1]. 110 | """ 111 | if reformulate: 112 | lam = beta.rvs(alpha+1, alpha) 113 | else: 114 | lam = beta.rvs(alpha, alpha) 115 | 116 | return lam 117 | 118 | 119 | def binarise_mask(mask, lam, in_shape, max_soft=0.0): 120 | """ Binarises a given low frequency image such that it has mean lambda. 121 | :param mask: Low frequency image, usually the result of `make_low_freq_image` 122 | :param lam: Mean value of final mask 123 | :param in_shape: Shape of inputs 124 | :param max_soft: Softening value between 0 and 0.5 which smooths hard edges in the mask. 125 | :return: 126 | """ 127 | idx = mask.reshape(-1).argsort()[::-1] 128 | mask = mask.reshape(-1) 129 | num = math.ceil(lam * mask.size) if random.random() > 0.5 else math.floor(lam * mask.size) # nosec # noqa 130 | 131 | eff_soft = max_soft 132 | if max_soft > lam or max_soft > (1-lam): 133 | eff_soft = min(lam, 1-lam) 134 | 135 | soft = int(mask.size * eff_soft) 136 | num_low = num - soft 137 | num_high = num + soft 138 | 139 | mask[idx[:num_high]] = 1 140 | mask[idx[num_low:]] = 0 141 | mask[idx[num_low:num_high]] = np.linspace(1, 0, (num_high - num_low)) 142 | 143 | mask = mask.reshape((1, *in_shape)) 144 | return mask 145 | 146 | 147 | def sample_mask(alpha, decay_power, shape, max_soft=0.0, reformulate=False): 148 | """ Samples a mean lambda from beta distribution parametrised by alpha, creates a low frequency image and binarises 149 | it based on this lambda 150 | :param alpha: Alpha value for beta distribution from which to sample mean of mask 151 | :param decay_power: Decay power for frequency decay prop 1/f**d 152 | :param shape: Shape of desired mask, list up to 3 dims 153 | :param max_soft: Softening value between 0 and 0.5 which smooths hard edges in the mask. 154 | :param reformulate: If True, uses the reformulation of [1]. 155 | """ 156 | if isinstance(shape, int): 157 | shape = (shape,) 158 | # Choose lambda 159 | lam = sample_lam(alpha, reformulate) 160 | 161 | # Make mask, get mean / std 162 | mask = make_low_freq_image(decay_power, shape) 163 | mask = binarise_mask(mask, lam, shape, max_soft) 164 | 165 | return lam, mask 166 | 167 | 168 | def sample_and_apply(x, alpha, decay_power, shape, max_soft=0.0, reformulate=False): 169 | """ 170 | :param x: Image batch on which to apply fmix of shape [b, c, shape*] 171 | :param alpha: Alpha value for beta distribution from which to sample mean of mask 172 | :param decay_power: Decay power for frequency decay prop 1/f**d 173 | :param shape: Shape of desired mask, list up to 3 dims 174 | :param max_soft: Softening value between 0 and 0.5 which smooths hard edges in the mask. 175 | :param reformulate: If True, uses the reformulation of [1]. 176 | :return: mixed input, permutation indices, lambda value of mix, 177 | """ 178 | lam, mask = sample_mask(alpha, decay_power, shape, max_soft, reformulate) 179 | index = np.random.permutation(x.shape[0]) 180 | 181 | x1, x2 = x * mask, x[index] * (1-mask) 182 | return x1+x2, index, lam 183 | 184 | 185 | class FMixBase: 186 | r""" FMix augmentation 187 | Args: 188 | decay_power (float): Decay power for frequency decay prop 1/f**d 189 | alpha (float): Alpha value for beta distribution from which to sample mean of mask 190 | size ([int] | [int, int] | [int, int, int]): Shape of desired mask, list up to 3 dims 191 | max_soft (float): Softening value between 0 and 0.5 which smooths hard edges in the mask. 192 | reformulate (bool): If True, uses the reformulation of [1]. 193 | """ 194 | 195 | def __init__(self, decay_power=3, alpha=1, size=(32, 32), max_soft=0.0, reformulate=False): 196 | super().__init__() 197 | self.decay_power = decay_power 198 | self.reformulate = reformulate 199 | self.size = size 200 | self.alpha = alpha 201 | self.max_soft = max_soft 202 | self.index = None 203 | self.lam = None 204 | 205 | def __call__(self, x): 206 | raise NotImplementedError 207 | 208 | def loss(self, *args, **kwargs): 209 | raise NotImplementedError 210 | -------------------------------------------------------------------------------- /torchreid/utils/ie_tools.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) 2021 Intel Corporation 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import logging as log 18 | import os 19 | 20 | import cv2 as cv 21 | import numpy as np 22 | from torchreid.models.common import ModelInterface 23 | 24 | 25 | class IEModel: 26 | """Class for inference of models in the Inference Engine format""" 27 | def __init__(self, exec_net, inputs_info, input_key, output_key, switch_rb=False): 28 | self.net = exec_net 29 | self.inputs_info = inputs_info 30 | self.input_key = input_key 31 | self.output_key = output_key 32 | self.reqs_ids = [] 33 | self.switch_rb = switch_rb 34 | 35 | def _preprocess(self, img): 36 | _, _, h, w = self.get_input_shape() 37 | if self.switch_rb: 38 | img = cv.cvtColor(img, cv.COLOR_BGR2RGB) 39 | img = np.expand_dims(cv.resize(img, (w, h)).transpose(2, 0, 1), axis=0) 40 | return img 41 | 42 | def forward(self, img): 43 | """Performs forward pass of the wrapped IE model""" 44 | res = self.net.infer(inputs={self.input_key: self._preprocess(img)}) 45 | return np.copy(res[self.output_key]) 46 | 47 | def forward_async(self, img): 48 | r_id = len(self.reqs_ids) 49 | self.net.start_async(request_id=r_id, 50 | inputs={self.input_key: self._preprocess(img)}) 51 | self.reqs_ids.append(r_id) 52 | 53 | def grab_all_async(self): 54 | outputs = [] 55 | for r_id in self.reqs_ids: 56 | self.net.requests[r_id].wait(-1) 57 | res = self.net.requests[r_id].output_blobs[self.output_key].buffer 58 | outputs.append(np.copy(res)) 59 | self.reqs_ids = [] 60 | return outputs 61 | 62 | def get_input_shape(self): 63 | """Returns an input shape of the wrapped IE model""" 64 | return self.inputs_info[self.input_key].input_data.shape 65 | 66 | 67 | def load_ie_model(ie, model_xml, device, plugin_dir, cpu_extension='', num_reqs=1, **kwargs): 68 | """Loads a model in the Inference Engine format""" 69 | # Plugin initialization for specified device and load extensions library if specified 70 | log.info(f"Initializing Inference Engine plugin for {device}") 71 | 72 | if cpu_extension and 'CPU' in device: 73 | ie.add_extension(cpu_extension, 'CPU') 74 | # Read IR 75 | log.info("Loading network") 76 | net = ie.read_network(model_xml, os.path.splitext(model_xml)[0] + ".bin") 77 | 78 | assert len(net.input_info) == 1 or len(net.input_info) == 2, \ 79 | "Supports topologies with only 1 or 2 inputs" 80 | assert len(net.outputs) == 1 or len(net.outputs) == 4 or len(net.outputs) == 5, \ 81 | "Supports topologies with only 1, 4 or 5 outputs" 82 | 83 | log.info("Preparing input blobs") 84 | input_blob = next(iter(net.input_info)) 85 | out_blob = next(iter(net.outputs)) 86 | net.batch_size = 1 87 | 88 | # Loading model to the plugin 89 | log.info("Loading model to the plugin") 90 | exec_net = ie.load_network(network=net, device_name=device, num_requests=num_reqs) 91 | model = IEModel(exec_net, net.input_info, input_blob, out_blob, **kwargs) 92 | return model 93 | 94 | 95 | class VectorCNN(ModelInterface): 96 | """Wrapper class for a nework returning a vector""" 97 | def __init__(self, ie, model_path, device='CPU', switch_rb=False, **kwargs): 98 | super().__init__(**kwargs) 99 | self.net = load_ie_model(ie, model_path, device, None, '', switch_rb=switch_rb) 100 | self.is_ie_model = True 101 | 102 | def forward(self, batch): 103 | """Performs forward of the underlying network on a given batch""" 104 | outputs = [self.net.forward(frame) for frame in batch] 105 | return outputs 106 | -------------------------------------------------------------------------------- /torchreid/utils/loggers.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (c) 2017 Tong Xiao 5 | # SPDX-License-Identifier: MIT 6 | # 7 | # Copyright (C) 2020-2021 Intel Corporation 8 | # SPDX-License-Identifier: Apache-2.0 9 | # 10 | 11 | # pylint: disable=consider-using-with 12 | 13 | from __future__ import absolute_import 14 | from datetime import datetime 15 | import os 16 | import os.path as osp 17 | import sys 18 | 19 | from .tools import mkdir_if_missing 20 | 21 | __all__ = ['Logger'] 22 | 23 | 24 | class Logger: 25 | """Writes console output to external text file. 26 | 27 | Imported from ``_ 28 | 29 | Args: 30 | fpath (str): directory to save logging file. 31 | 32 | Examples:: 33 | >>> import sys 34 | >>> import os 35 | >>> import os.path as osp 36 | >>> from torchreid.utils import Logger 37 | >>> save_dir = 'log/resnet50-softmax-market1501' 38 | >>> log_name = 'train.log' 39 | >>> sys.stdout = Logger(osp.join(args.save_dir, log_name)) 40 | """ 41 | 42 | def __init__(self, fpath=None): 43 | self.console = sys.stdout 44 | self.file = None 45 | self.was_cr = True 46 | if fpath is not None: 47 | mkdir_if_missing(osp.dirname(fpath)) 48 | self.file = open(fpath, 'w') 49 | 50 | def __del__(self): 51 | self.close() 52 | 53 | def __enter__(self): 54 | pass 55 | 56 | def __exit__(self, *args): 57 | self.close() 58 | 59 | def write(self, msg): 60 | if self.was_cr: 61 | timestamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S|') 62 | else: 63 | timestamp = '' 64 | self.was_cr = msg.endswith('\n') 65 | 66 | self.console.write(timestamp + msg) 67 | if self.file is not None: 68 | self.file.write(timestamp + msg) 69 | 70 | def flush(self): 71 | self.console.flush() 72 | if self.file is not None: 73 | self.file.flush() 74 | os.fsync(self.file.fileno()) 75 | 76 | def close(self): 77 | self.console.close() 78 | if self.file is not None: 79 | self.file.close() 80 | -------------------------------------------------------------------------------- /torchreid/utils/tools.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (c) 2018 davidtvs 5 | # SPDX-License-Identifier: MIT 6 | # 7 | # Copyright (c) 2018 Facebook 8 | # SPDX-License-Identifier: MIT 9 | # 10 | # Copyright (C) 2020-2021 Intel Corporation 11 | # SPDX-License-Identifier: Apache-2.0 12 | # 13 | 14 | from __future__ import absolute_import, division, print_function 15 | import copy 16 | import errno 17 | import os 18 | import os.path as osp 19 | import random 20 | import subprocess 21 | 22 | import numpy as np 23 | import torch 24 | import cv2 as cv 25 | 26 | __all__ = [ 27 | 'mkdir_if_missing', 'check_isfile', 'set_random_seed', "worker_init_fn", 28 | 'read_image', 'get_model_attr', 'StateCacher', 'random_image', 'EvalModeSetter', 29 | 'get_git_revision', 'set_model_attr' 30 | ] 31 | 32 | def get_git_revision(): 33 | path = os.path.abspath(os.path.dirname(__file__)) 34 | sha_message = ['git', 'rev-parse', 'HEAD'] 35 | head_message = sha_message[:2] + ['--abbrev-ref'] + sha_message[2:] 36 | return (subprocess.check_output(sha_message, cwd=path).decode('ascii').strip(), 37 | subprocess.check_output(head_message, cwd=path).decode('ascii').strip()) 38 | 39 | 40 | def mkdir_if_missing(dirname): 41 | """Creates dirname if it is missing.""" 42 | if not osp.exists(dirname): 43 | try: 44 | os.makedirs(dirname) 45 | except OSError as e: 46 | if e.errno != errno.EEXIST: 47 | raise 48 | 49 | def check_isfile(fpath): 50 | """Checks if the given path is a file. 51 | 52 | Args: 53 | fpath (str): file path. 54 | 55 | Returns: 56 | bool 57 | """ 58 | isfile = osp.isfile(fpath) 59 | return isfile 60 | 61 | def set_random_seed(seed, deterministic=False): 62 | torch.manual_seed(seed) 63 | torch.cuda.manual_seed_all(seed) 64 | if deterministic: 65 | torch.backends.cudnn.deterministic = True 66 | torch.backends.cudnn.benchmark = False 67 | np.random.seed(seed) 68 | random.seed(seed) 69 | os.environ['PYTHONHASHSEED'] = str(seed) 70 | 71 | def worker_init_fn(worker_id): 72 | np.random.seed(np.random.get_state()[1][0] + worker_id) 73 | random.seed(random.getstate()[1][0] + worker_id) 74 | 75 | 76 | def read_image(path, grayscale=False): 77 | """Reads image from path using ``Open CV``. 78 | 79 | Args: 80 | path (str): path to an image. 81 | grayscale (bool): load grayscale image 82 | 83 | Returns: 84 | Numpy image 85 | """ 86 | 87 | got_img = False 88 | if not osp.exists(path): 89 | raise IOError(f'"{path}" does not exist') 90 | 91 | while not got_img: 92 | try: 93 | img = cv.cvtColor(cv.imread(path, cv.IMREAD_COLOR), cv.COLOR_BGR2RGB) 94 | got_img = True 95 | except IOError: 96 | print(f'IOError occurred when reading "{path}".') 97 | 98 | return img 99 | 100 | 101 | def random_image(height, width): 102 | input_size = (height, width, 3) 103 | img = np.random.rand(*input_size).astype(np.float32) 104 | img = np.uint8(img * 255) 105 | 106 | return img 107 | 108 | 109 | def get_model_attr(model, attr): 110 | if hasattr(model, 'module'): 111 | model = model.module 112 | return getattr(model, attr) 113 | 114 | 115 | def set_model_attr(model, attr, value): 116 | if hasattr(model, 'module'): 117 | model = model.module 118 | if hasattr(model, 'nncf_module'): 119 | setattr(model.nncf_module, attr, value) 120 | setattr(model, attr, value) 121 | 122 | 123 | class StateCacher: 124 | def __init__(self, in_memory, cache_dir=None): 125 | self.in_memory = in_memory 126 | self.cache_dir = cache_dir 127 | 128 | if self.cache_dir is None: 129 | import tempfile 130 | 131 | self.cache_dir = tempfile.gettempdir() 132 | else: 133 | if not os.path.isdir(self.cache_dir): 134 | raise ValueError("Given `cache_dir` is not a valid directory.") 135 | 136 | self.cached = {} 137 | 138 | def store(self, key, state_dict): 139 | if self.in_memory: 140 | self.cached.update({key: copy.deepcopy(state_dict)}) 141 | else: 142 | fn = os.path.join(self.cache_dir, f"state_{key}_{id(self)}.pt") 143 | self.cached.update({key: fn}) 144 | torch.save(state_dict, fn) 145 | 146 | def retrieve(self, key): 147 | if key not in self.cached: 148 | raise KeyError(f"Target {key} was not cached.") 149 | 150 | if self.in_memory: 151 | return self.cached.get(key) 152 | 153 | fn = self.cached.get(key) 154 | if not os.path.exists(fn): 155 | raise RuntimeError( 156 | f"Failed to load state in {fn}. File doesn't exist anymore." 157 | ) 158 | state_dict = torch.load(fn, map_location=lambda storage, location: storage) 159 | return state_dict 160 | 161 | def __del__(self): 162 | """Check whether there are unused cached files existing in `cache_dir` before 163 | this instance being destroyed.""" 164 | 165 | if self.in_memory: 166 | return 167 | 168 | for _, v in self.cached.items(): 169 | if os.path.exists(v): 170 | os.remove(v) 171 | 172 | 173 | class EvalModeSetter: 174 | def __init__(self, module, m_type): 175 | self.modules = module 176 | if not isinstance(self.modules, (tuple, list)): 177 | self.modules = [self.modules] 178 | 179 | self.modes_storage = [{} for _ in range(len(self.modules))] 180 | 181 | self.m_types = m_type 182 | if not isinstance(self.m_types, (tuple, list)): 183 | self.m_types = [self.m_types] 184 | 185 | def __enter__(self): 186 | for module_id, module in enumerate(self.modules): 187 | modes_storage = self.modes_storage[module_id] 188 | 189 | for child_name, child_module in module.named_modules(): 190 | matched = any(isinstance(child_module, m_type) for m_type in self.m_types) 191 | if matched: 192 | modes_storage[child_name] = child_module.training 193 | child_module.train(mode=False) 194 | 195 | return self 196 | 197 | def __exit__(self, exc_type, exc_val, exc_tb): 198 | for module_id, module in enumerate(self.modules): 199 | modes_storage = self.modes_storage[module_id] 200 | 201 | for child_name, child_module in module.named_modules(): 202 | if child_name in modes_storage: 203 | child_module.train(mode=modes_storage[child_name]) 204 | -------------------------------------------------------------------------------- /torchreid/version.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2021 Kaiyang Zhou 2 | # SPDX-License-Identifier: MIT 3 | # 4 | # Copyright (C) 2020-2021 Intel Corporation 5 | # SPDX-License-Identifier: Apache-2.0 6 | # 7 | 8 | __version__ = '1.2.3' 9 | --------------------------------------------------------------------------------