├── .gitignore ├── .style.yapf ├── LICENSE ├── README.md ├── benchmarks ├── detection │ ├── README.md │ ├── configs │ │ ├── Base-Keypoint-RCNN-FPN.yaml │ │ ├── Base-RCNN-C4-BN.yaml │ │ ├── Base-RCNN-FPN.yaml │ │ ├── Base-RetinaNet.yaml │ │ ├── Cityscapes │ │ │ ├── mask_rcnn_R_50_FPN.yaml │ │ │ └── mask_rcnn_R_50_FPN_moco.yaml │ │ ├── coco_R_50_C4_1x.yaml │ │ ├── coco_R_50_C4_1x_moco.yaml │ │ ├── coco_R_50_C4_2x.yaml │ │ ├── coco_R_50_C4_2x_moco.yaml │ │ ├── coco_R_50_FPN_1x.yaml │ │ ├── coco_R_50_FPN_1x_moco.yaml │ │ ├── coco_R_50_FPN_2x.yaml │ │ ├── coco_R_50_FPN_2x_moco.yaml │ │ ├── coco_R_50_RetinaNet_1x.yaml │ │ ├── coco_R_50_RetinaNet_1x_moco.yaml │ │ ├── coco_R_50_RetinaNet_2x.yaml │ │ ├── coco_R_50_RetinaNet_2x_moco.yaml │ │ ├── keypoint_rcnn_R_50_FPN_2x.yaml │ │ ├── keypoint_rcnn_R_50_FPN_2x_moco.yaml │ │ ├── pascal_voc_R_50_C4_24k.yaml │ │ └── pascal_voc_R_50_C4_24k_moco.yaml │ ├── convert-pretrain-to-detectron2.py │ ├── run.sh │ └── train_net.py ├── dist_test_svm_epoch.sh ├── dist_test_svm_pretrain.sh ├── dist_train_linear.sh ├── dist_train_semi.sh ├── extract_info │ └── voc07.py ├── srun_test_svm_epoch.sh ├── srun_test_svm_pretrain.sh ├── srun_train_linear.sh ├── srun_train_semi.sh └── svm_tools │ ├── aggregate_low_shot_svm_stats.py │ ├── eval_svm_full.sh │ ├── eval_svm_lowshot.sh │ ├── svm_helper.py │ ├── test_svm.py │ ├── test_svm_low_shot.py │ ├── train_svm_kfold.py │ ├── train_svm_kfold_parallel.py │ ├── train_svm_low_shot.py │ └── train_svm_low_shot_parallel.py ├── configs ├── base.py ├── benchmarks │ ├── linear_classification │ │ ├── imagenet │ │ │ ├── r50_last.py │ │ │ ├── r50_last_sobel.py │ │ │ ├── r50_multihead.py │ │ │ └── r50_multihead_sobel.py │ │ └── places205 │ │ │ ├── r50_multihead.py │ │ │ └── r50_multihead_sobel.py │ └── semi_classification │ │ ├── imagenet_10percent │ │ ├── base.py │ │ ├── r50_lr0_001_head1.py │ │ ├── r50_lr0_001_head10.py │ │ ├── r50_lr0_001_head100.py │ │ ├── r50_lr0_01_head1.py │ │ ├── r50_lr0_01_head10.py │ │ ├── r50_lr0_01_head100.py │ │ ├── r50_lr0_01_head1_sobel.py │ │ ├── r50_lr0_1_head1.py │ │ ├── r50_lr0_1_head10.py │ │ └── r50_lr0_1_head100.py │ │ └── imagenet_1percent │ │ ├── base.py │ │ ├── r50_lr0_001_head1.py │ │ ├── r50_lr0_001_head10.py │ │ ├── r50_lr0_001_head100.py │ │ ├── r50_lr0_01_head1.py │ │ ├── r50_lr0_01_head10.py │ │ ├── r50_lr0_01_head100.py │ │ ├── r50_lr0_01_head1_sobel.py │ │ ├── r50_lr0_1_head1.py │ │ ├── r50_lr0_1_head10.py │ │ └── r50_lr0_1_head100.py ├── classification │ ├── cifar10 │ │ └── r50.py │ └── imagenet │ │ └── r50.py └── selfsup │ ├── byol │ ├── r50_bs2048_accumulate2_ep200_fp16.py │ ├── r50_bs256_accumulate16_ep300.py │ └── r50_bs256_ep200.py │ ├── deepcluster │ ├── r50.py │ └── r50_withoutsobel.py │ ├── moco │ ├── r50_v1.py │ ├── r50_v1_fp16.py │ ├── r50_v2.py │ └── r50_v2_simclr_neck.py │ ├── npid │ ├── r50.py │ └── r50_ensure_neg.py │ ├── odc │ └── r50_v1.py │ ├── regionCL-M │ └── regionCLM-r50.py │ ├── relative_loc │ └── r50.py │ ├── rotation_pred │ ├── r50.py │ └── r50_fp16.py │ └── simclr │ ├── r50_bs256_ep200.py │ ├── r50_bs256_ep200_mocov2_neck.py │ └── r50_bs512_ep200.py ├── demo └── NetworkStructure.png ├── docs ├── CHANGELOG.md ├── GETTING_STARTED.md ├── INSTALL.md ├── MODEL_ZOO.md └── relation.jpg ├── logs └── RegionCLM │ ├── CityScape-AP-RegionCLM.log │ ├── CityScape-UperNet-RegionCLM-40k.log │ ├── CityScape-UperNet-RegionCLM-80k.log │ ├── Linear-RegionCLM.log │ ├── MaskRCNN-C4-RegionCLM-1x.log │ ├── MaskRCNN-C4-RegionCLM-2x.log │ ├── MaskRCNN-FPN-RegionCLM-1x.log │ ├── MaskRCNN-FPN-RegionCLM-2x.log │ ├── RegionCLM.log │ ├── RetinaNet-RegionCLM-1x.log │ └── RetinaNet-RegionCLM-2x.log ├── openselfsup ├── __init__.py ├── apis │ ├── __init__.py │ └── train.py ├── datasets │ ├── __init__.py │ ├── base.py │ ├── builder.py │ ├── byol.py │ ├── classification.py │ ├── contrastive.py │ ├── data_sources │ │ ├── __init__.py │ │ ├── cifar.py │ │ ├── image_list.py │ │ ├── imagenet.py │ │ ├── places205.py │ │ └── utils.py │ ├── dataset_wrappers.py │ ├── deepcluster.py │ ├── extraction.py │ ├── loader │ │ ├── __init__.py │ │ ├── build_loader.py │ │ └── sampler.py │ ├── npid.py │ ├── pipelines │ │ ├── __init__.py │ │ └── transforms.py │ ├── registry.py │ ├── relative_loc.py │ ├── rotation_pred.py │ └── utils.py ├── hooks │ ├── __init__.py │ ├── builder.py │ ├── byol_hook.py │ ├── deepcluster_hook.py │ ├── extractor.py │ ├── odc_hook.py │ ├── optimizer_hook.py │ ├── registry.py │ └── validate_hook.py ├── models │ ├── RegionCLM.py │ ├── __init__.py │ ├── backbones │ │ ├── __init__.py │ │ ├── resnet.py │ │ └── resnext.py │ ├── builder.py │ ├── byol.py │ ├── classification.py │ ├── deepcluster.py │ ├── heads │ │ ├── __init__.py │ │ ├── cls_head.py │ │ ├── contrastive_head.py │ │ ├── latent_pred_head.py │ │ └── multi_cls_head.py │ ├── memories │ │ ├── __init__.py │ │ ├── odc_memory.py │ │ └── simple_memory.py │ ├── moco.py │ ├── necks.py │ ├── npid.py │ ├── odc.py │ ├── registry.py │ ├── relative_loc.py │ ├── rotation_pred.py │ ├── simclr.py │ └── utils │ │ ├── __init__.py │ │ ├── accuracy.py │ │ ├── conv_module.py │ │ ├── conv_ws.py │ │ ├── gather_layer.py │ │ ├── multi_pooling.py │ │ ├── norm.py │ │ ├── scale.py │ │ └── sobel.py ├── third_party │ └── clustering.py └── utils │ ├── __init__.py │ ├── alias_multinomial.py │ ├── collect.py │ ├── collect_env.py │ ├── config_tools.py │ ├── contextmanagers.py │ ├── flops_counter.py │ ├── gather.py │ ├── logger.py │ ├── misc.py │ ├── optimizers.py │ ├── profiling.py │ └── registry.py ├── requirements.txt ├── requirements ├── runtime.txt └── tests.txt ├── setup.py └── tools ├── count_parameters.py ├── dist_extract.sh ├── dist_test.sh ├── dist_train.sh ├── extract.py ├── extract_backbone_weights.py ├── kill.sh ├── prepare_data ├── convert_subset.py ├── create_voc_data_files.py ├── create_voc_low_shot_challenge_samples.py └── prepare_voc07_cls.sh ├── publish_model.py ├── single_train.sh ├── srun_extract.sh ├── srun_test.sh ├── srun_train.sh ├── test.py ├── train.py └── upgrade_models.py /.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 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | # *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | openselfsup/version.py 107 | version.py 108 | data 109 | .vscode 110 | .idea 111 | 112 | # custom 113 | *.pkl 114 | *.pkl.json 115 | *.log.json 116 | work_dirs/ 117 | pretrains 118 | 119 | # Pytorch 120 | *.pth 121 | 122 | *.swp 123 | source.sh 124 | tensorboard.sh 125 | .DS_Store 126 | replace.sh 127 | benchmarks/detection/datasets 128 | benchmarks/detection/output 129 | -------------------------------------------------------------------------------- /.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 | -------------------------------------------------------------------------------- /benchmarks/detection/README.md: -------------------------------------------------------------------------------- 1 | 2 | ## Transferring to Detection 3 | 4 | We follow the evaluation setting in MoCo when trasferring to object detection. 5 | 6 | ### Instruction 7 | 8 | 1. Install [detectron2](https://github.com/facebookresearch/detectron2/blob/master/INSTALL.md). 9 | 10 | 1. Put dataset under "benchmarks/detection/datasets" directory, 11 | following the [directory structure](https://github.com/facebookresearch/detectron2/tree/master/datasets) 12 | requried by detectron2. 13 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/Base-Keypoint-RCNN-FPN.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "Base-RCNN-FPN.yaml" 2 | MODEL: 3 | KEYPOINT_ON: True 4 | ROI_HEADS: 5 | NUM_CLASSES: 1 6 | ROI_BOX_HEAD: 7 | SMOOTH_L1_BETA: 0.5 # Keypoint AP degrades (though box AP improves) when using plain L1 loss 8 | RPN: 9 | # Detectron1 uses 2000 proposals per-batch, but this option is per-image in detectron2. 10 | # 1000 proposals per-image is found to hurt box AP. 11 | # Therefore we increase it to 1500 per-image. 12 | POST_NMS_TOPK_TRAIN: 1500 13 | DATASETS: 14 | TRAIN: ("keypoints_coco_2017_train",) 15 | TEST: ("keypoints_coco_2017_val",) 16 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/Base-RCNN-C4-BN.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | META_ARCHITECTURE: "GeneralizedRCNN" 3 | RPN: 4 | PRE_NMS_TOPK_TEST: 6000 5 | POST_NMS_TOPK_TEST: 1000 6 | ROI_HEADS: 7 | NAME: "Res5ROIHeadsExtraNorm" 8 | BACKBONE: 9 | FREEZE_AT: 0 10 | RESNETS: 11 | NORM: "SyncBN" 12 | TEST: 13 | PRECISE_BN: 14 | ENABLED: True 15 | SOLVER: 16 | IMS_PER_BATCH: 16 17 | BASE_LR: 0.02 18 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/Base-RCNN-FPN.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | META_ARCHITECTURE: "GeneralizedRCNN" 3 | BACKBONE: 4 | NAME: "build_resnet_fpn_backbone" 5 | RESNETS: 6 | OUT_FEATURES: ["res2", "res3", "res4", "res5"] 7 | FPN: 8 | IN_FEATURES: ["res2", "res3", "res4", "res5"] 9 | ANCHOR_GENERATOR: 10 | SIZES: [[32], [64], [128], [256], [512]] # One size for each in feature map 11 | ASPECT_RATIOS: [[0.5, 1.0, 2.0]] # Three aspect ratios (same for all in feature maps) 12 | RPN: 13 | IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"] 14 | PRE_NMS_TOPK_TRAIN: 2000 # Per FPN level 15 | PRE_NMS_TOPK_TEST: 1000 # Per FPN level 16 | # Detectron1 uses 2000 proposals per-batch, 17 | # (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue) 18 | # which is approximately 1000 proposals per-image since the default batch size for FPN is 2. 19 | POST_NMS_TOPK_TRAIN: 1000 20 | POST_NMS_TOPK_TEST: 1000 21 | ROI_HEADS: 22 | NAME: "StandardROIHeads" 23 | IN_FEATURES: ["p2", "p3", "p4", "p5"] 24 | ROI_BOX_HEAD: 25 | NAME: "FastRCNNConvFCHead" 26 | NUM_FC: 2 27 | POOLER_RESOLUTION: 7 28 | ROI_MASK_HEAD: 29 | NAME: "MaskRCNNConvUpsampleHead" 30 | NUM_CONV: 4 31 | POOLER_RESOLUTION: 14 32 | DATASETS: 33 | TRAIN: ("coco_2017_train",) 34 | TEST: ("coco_2017_val",) 35 | SOLVER: 36 | IMS_PER_BATCH: 16 37 | BASE_LR: 0.02 38 | STEPS: (60000, 80000) 39 | MAX_ITER: 90000 40 | INPUT: 41 | MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800) 42 | VERSION: 2 -------------------------------------------------------------------------------- /benchmarks/detection/configs/Base-RetinaNet.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | META_ARCHITECTURE: "RetinaNet" 3 | BACKBONE: 4 | NAME: "build_retinanet_resnet_fpn_backbone" 5 | RESNETS: 6 | OUT_FEATURES: ["res3", "res4", "res5"] 7 | ANCHOR_GENERATOR: 8 | SIZES: !!python/object/apply:eval ["[[x, x * 2**(1.0/3), x * 2**(2.0/3) ] for x in [32, 64, 128, 256, 512 ]]"] 9 | FPN: 10 | IN_FEATURES: ["res3", "res4", "res5"] 11 | RETINANET: 12 | IOU_THRESHOLDS: [0.4, 0.5] 13 | IOU_LABELS: [0, -1, 1] 14 | SMOOTH_L1_LOSS_BETA: 0.0 15 | DATASETS: 16 | TRAIN: ("coco_2017_train",) 17 | TEST: ("coco_2017_val",) 18 | SOLVER: 19 | IMS_PER_BATCH: 16 20 | BASE_LR: 0.01 # Note that RetinaNet uses a different default learning rate 21 | STEPS: (60000, 80000) 22 | MAX_ITER: 90000 23 | INPUT: 24 | MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800) 25 | VERSION: 2 -------------------------------------------------------------------------------- /benchmarks/detection/configs/Cityscapes/mask_rcnn_R_50_FPN.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "../Base-RCNN-FPN.yaml" 2 | MODEL: 3 | WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl" 4 | MASK_ON: True 5 | ROI_HEADS: 6 | NUM_CLASSES: 8 7 | BACKBONE: 8 | FREEZE_AT: 0 9 | RESNETS: 10 | DEPTH: 50 11 | NORM: "SyncBN" 12 | FPN: 13 | NORM: "SyncBN" 14 | INPUT: 15 | MIN_SIZE_TRAIN: (800, 832, 864, 896, 928, 960, 992, 1024) 16 | MIN_SIZE_TRAIN_SAMPLING: "choice" 17 | MIN_SIZE_TEST: 1024 18 | MAX_SIZE_TRAIN: 2048 19 | MAX_SIZE_TEST: 2048 20 | DATASETS: 21 | TRAIN: ("cityscapes_fine_instance_seg_train",) 22 | TEST: ("cityscapes_fine_instance_seg_val",) 23 | SOLVER: 24 | BASE_LR: 0.01 25 | STEPS: (18000,) 26 | MAX_ITER: 24000 27 | IMS_PER_BATCH: 8 28 | TEST: 29 | PRECISE_BN: 30 | ENABLED: True -------------------------------------------------------------------------------- /benchmarks/detection/configs/Cityscapes/mask_rcnn_R_50_FPN_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "mask_rcnn_R_50_FPN.yaml" 2 | MODEL: 3 | PIXEL_MEAN: [123.675, 116.280, 103.530] 4 | PIXEL_STD: [58.395, 57.120, 57.375] 5 | WEIGHTS: "See Instructions" 6 | RESNETS: 7 | STRIDE_IN_1X1: False 8 | INPUT: 9 | FORMAT: "RGB" -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_C4_1x.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_C4_2x.yaml" 2 | SOLVER: 3 | STEPS: (60000, 80000) 4 | MAX_ITER: 90000 5 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_C4_1x_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_C4_2x_moco.yaml" 2 | SOLVER: 3 | STEPS: (60000, 80000) 4 | MAX_ITER: 90000 5 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_C4_2x.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "Base-RCNN-C4-BN.yaml" 2 | MODEL: 3 | MASK_ON: True 4 | WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl" 5 | INPUT: 6 | MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800) 7 | MIN_SIZE_TEST: 800 8 | DATASETS: 9 | TRAIN: ("coco_2017_train",) 10 | TEST: ("coco_2017_val",) 11 | SOLVER: 12 | STEPS: (120000, 160000) 13 | MAX_ITER: 180000 14 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_C4_2x_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_C4_2x.yaml" 2 | MODEL: 3 | PIXEL_MEAN: [123.675, 116.280, 103.530] 4 | PIXEL_STD: [58.395, 57.120, 57.375] 5 | WEIGHTS: "See Instructions" 6 | RESNETS: 7 | STRIDE_IN_1X1: False 8 | INPUT: 9 | MAX_SIZE_TRAIN: 1200 10 | FORMAT: "RGB" 11 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_FPN_1x.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "Base-RCNN-FPN.yaml" 2 | MODEL: 3 | MASK_ON: True 4 | WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl" 5 | BACKBONE: 6 | FREEZE_AT: 0 7 | RESNETS: 8 | DEPTH: 50 9 | NORM: "SyncBN" 10 | FPN: 11 | NORM: "SyncBN" 12 | TEST: 13 | PRECISE_BN: 14 | ENABLED: True 15 | SOLVER: 16 | STEPS: (60000, 80000) 17 | MAX_ITER: 90000 -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_FPN_1x_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_FPN_1x.yaml" 2 | MODEL: 3 | PIXEL_MEAN: [123.675, 116.280, 103.530] 4 | PIXEL_STD: [58.395, 57.120, 57.375] 5 | WEIGHTS: "See Instructions" 6 | RESNETS: 7 | STRIDE_IN_1X1: False 8 | INPUT: 9 | FORMAT: "RGB" -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_FPN_2x.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_FPN_1x.yaml" 2 | SOLVER: 3 | STEPS: (120000, 160000) 4 | MAX_ITER: 180000 -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_FPN_2x_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_FPN_1x_moco.yaml" 2 | SOLVER: 3 | STEPS: (120000, 160000) 4 | MAX_ITER: 180000 -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_RetinaNet_1x.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "Base-RetinaNet.yaml" 2 | MODEL: 3 | WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl" 4 | BACKBONE: 5 | FREEZE_AT: 0 6 | RESNETS: 7 | DEPTH: 50 8 | NORM: "SyncBN" 9 | FPN: 10 | NORM: "SyncBN" 11 | TEST: 12 | PRECISE_BN: 13 | ENABLED: True -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_RetinaNet_1x_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_RetinaNet_1x.yaml" 2 | MODEL: 3 | PIXEL_MEAN: [123.675, 116.280, 103.530] 4 | PIXEL_STD: [58.395, 57.120, 57.375] 5 | WEIGHTS: "See Instructions" 6 | RESNETS: 7 | STRIDE_IN_1X1: False 8 | INPUT: 9 | FORMAT: "RGB" -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_RetinaNet_2x.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_RetinaNet_1x.yaml" 2 | SOLVER: 3 | STEPS: (120000, 160000) 4 | MAX_ITER: 180000 -------------------------------------------------------------------------------- /benchmarks/detection/configs/coco_R_50_RetinaNet_2x_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "coco_R_50_RetinaNet_1x_moco.yaml" 2 | SOLVER: 3 | STEPS: (120000, 160000) 4 | MAX_ITER: 180000 -------------------------------------------------------------------------------- /benchmarks/detection/configs/keypoint_rcnn_R_50_FPN_2x.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "Base-Keypoint-RCNN-FPN.yaml" 2 | MODEL: 3 | WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl" 4 | BACKBONE: 5 | FREEZE_AT: 0 6 | RESNETS: 7 | DEPTH: 50 8 | NORM: "SyncBN" 9 | FPN: 10 | NORM: "SyncBN" 11 | TEST: 12 | PRECISE_BN: 13 | ENABLED: True 14 | SOLVER: 15 | STEPS: (120000, 160000) 16 | MAX_ITER: 180000 -------------------------------------------------------------------------------- /benchmarks/detection/configs/keypoint_rcnn_R_50_FPN_2x_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "keypoint_rcnn_R_50_FPN_2x.yaml" 2 | MODEL: 3 | PIXEL_MEAN: [123.675, 116.280, 103.530] 4 | PIXEL_STD: [58.395, 57.120, 57.375] 5 | WEIGHTS: "See Instructions" 6 | RESNETS: 7 | STRIDE_IN_1X1: False 8 | INPUT: 9 | FORMAT: "RGB" -------------------------------------------------------------------------------- /benchmarks/detection/configs/pascal_voc_R_50_C4_24k.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "Base-RCNN-C4-BN.yaml" 2 | MODEL: 3 | MASK_ON: False 4 | WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl" 5 | ROI_HEADS: 6 | NUM_CLASSES: 20 7 | INPUT: 8 | MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800) 9 | MIN_SIZE_TEST: 800 10 | DATASETS: 11 | TRAIN: ('voc_2007_trainval', 'voc_2012_trainval') 12 | TEST: ('voc_2007_test',) 13 | SOLVER: 14 | STEPS: (18000, 22000) 15 | MAX_ITER: 24000 16 | WARMUP_ITERS: 100 17 | -------------------------------------------------------------------------------- /benchmarks/detection/configs/pascal_voc_R_50_C4_24k_moco.yaml: -------------------------------------------------------------------------------- 1 | _BASE_: "pascal_voc_R_50_C4_24k.yaml" 2 | MODEL: 3 | PIXEL_MEAN: [123.675, 116.280, 103.530] 4 | PIXEL_STD: [58.395, 57.120, 57.375] 5 | WEIGHTS: "See Instructions" 6 | RESNETS: 7 | STRIDE_IN_1X1: False 8 | INPUT: 9 | FORMAT: "RGB" 10 | -------------------------------------------------------------------------------- /benchmarks/detection/convert-pretrain-to-detectron2.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 3 | 4 | import pickle as pkl 5 | import sys 6 | import torch 7 | 8 | if __name__ == "__main__": 9 | input = sys.argv[1] 10 | 11 | obj = torch.load(input, map_location="cpu") 12 | obj = obj["state_dict"] 13 | 14 | newmodel = {} 15 | for k, v in obj.items(): 16 | old_k = k 17 | if "layer" not in k: 18 | k = "stem." + k 19 | for t in [1, 2, 3, 4]: 20 | k = k.replace("layer{}".format(t), "res{}".format(t + 1)) 21 | for t in [1, 2, 3]: 22 | k = k.replace("bn{}".format(t), "conv{}.norm".format(t)) 23 | k = k.replace("downsample.0", "shortcut") 24 | k = k.replace("downsample.1", "shortcut.norm") 25 | print(old_k, "->", k) 26 | newmodel[k] = v.numpy() 27 | 28 | res = { 29 | "model": newmodel, 30 | "__author__": "OpenSelfSup", 31 | "matching_heuristics": True 32 | } 33 | 34 | assert sys.argv[2].endswith('.pkl') 35 | with open(sys.argv[2], "wb") as f: 36 | pkl.dump(res, f) 37 | -------------------------------------------------------------------------------- /benchmarks/detection/run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | DET_CFG=$1 3 | WEIGHTS=$2 4 | 5 | python $(dirname "$0")/train_net.py --config-file $DET_CFG \ 6 | --num-gpus 8 MODEL.WEIGHTS $WEIGHTS 7 | -------------------------------------------------------------------------------- /benchmarks/detection/train_net.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 3 | 4 | import os 5 | 6 | from detectron2.checkpoint import DetectionCheckpointer 7 | from detectron2.config import get_cfg 8 | from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, launch 9 | from detectron2.evaluation import COCOEvaluator, PascalVOCDetectionEvaluator 10 | from detectron2.layers import get_norm 11 | from detectron2.modeling.roi_heads import ROI_HEADS_REGISTRY, Res5ROIHeads 12 | 13 | 14 | @ROI_HEADS_REGISTRY.register() 15 | class Res5ROIHeadsExtraNorm(Res5ROIHeads): 16 | """ 17 | As described in the MOCO paper, there is an extra BN layer 18 | following the res5 stage. 19 | """ 20 | 21 | def _build_res5_block(self, cfg): 22 | seq, out_channels = super()._build_res5_block(cfg) 23 | norm = cfg.MODEL.RESNETS.NORM 24 | norm = get_norm(norm, out_channels) 25 | seq.add_module("norm", norm) 26 | return seq, out_channels 27 | 28 | 29 | class Trainer(DefaultTrainer): 30 | 31 | @classmethod 32 | def build_evaluator(cls, cfg, dataset_name, output_folder=None): 33 | if output_folder is None: 34 | output_folder = os.path.join(cfg.OUTPUT_DIR, "inference") 35 | if "coco" in dataset_name: 36 | return COCOEvaluator(dataset_name, cfg, True, output_folder) 37 | else: 38 | assert "voc" in dataset_name 39 | return PascalVOCDetectionEvaluator(dataset_name) 40 | 41 | 42 | def setup(args): 43 | cfg = get_cfg() 44 | cfg.merge_from_file(args.config_file) 45 | cfg.merge_from_list(args.opts) 46 | cfg.freeze() 47 | default_setup(cfg, args) 48 | return cfg 49 | 50 | 51 | def main(args): 52 | cfg = setup(args) 53 | 54 | if args.eval_only: 55 | model = Trainer.build_model(cfg) 56 | DetectionCheckpointer( 57 | model, save_dir=cfg.OUTPUT_DIR).resume_or_load( 58 | cfg.MODEL.WEIGHTS, resume=args.resume) 59 | res = Trainer.test(cfg, model) 60 | return res 61 | 62 | trainer = Trainer(cfg) 63 | trainer.resume_or_load(resume=args.resume) 64 | return trainer.train() 65 | 66 | 67 | if __name__ == "__main__": 68 | args = default_argument_parser().parse_args() 69 | print("Command Line Args:", args) 70 | launch( 71 | main, 72 | args.num_gpus, 73 | num_machines=args.num_machines, 74 | machine_rank=args.machine_rank, 75 | dist_url=args.dist_url, 76 | args=(args, ), 77 | ) 78 | -------------------------------------------------------------------------------- /benchmarks/dist_test_svm_epoch.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | set -x 4 | 5 | CFG=$1 6 | EPOCH=$2 7 | FEAT_LIST=$3 # e.g.: "feat5", "feat4 feat5". If leave empty, the default is "feat5" 8 | GPUS=${4:-8} 9 | WORK_DIR=$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/ 10 | 11 | if [ "$CFG" == "" ] || [ "$EPOCH" == "" ]; then 12 | echo "ERROR: Missing arguments." 13 | exit 14 | fi 15 | 16 | if [ ! -f $WORK_DIR/epoch_${EPOCH}.pth ]; then 17 | echo "ERROR: File not exist: $WORK_DIR/epoch_${EPOCH}.pth" 18 | exit 19 | fi 20 | 21 | mkdir -p $WORK_DIR/logs 22 | echo "Testing checkpoint: $WORK_DIR/epoch_${EPOCH}.pth" 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 23 | 24 | bash tools/dist_extract.sh $CFG $GPUS $WORK_DIR --checkpoint $WORK_DIR/epoch_${EPOCH}.pth 25 | 26 | bash benchmarks/svm_tools/eval_svm_full.sh $WORK_DIR "$FEAT_LIST" 27 | 28 | bash benchmarks/svm_tools/eval_svm_lowshot.sh $WORK_DIR "$FEAT_LIST" 29 | -------------------------------------------------------------------------------- /benchmarks/dist_test_svm_pretrain.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | set -x 4 | 5 | CFG=$1 6 | PRETRAIN=$2 # pretrained model or "random" (random init) 7 | FEAT_LIST=$3 # e.g.: "feat5", "feat4 feat5". If leave empty, the default is "feat5" 8 | GPUS=${4:-8} 9 | WORK_DIR="$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/$(echo $PRETRAIN | rev | cut -d/ -f 1 | rev)" 10 | 11 | if [ "$CFG" == "" ] || [ "$PRETRAIN" == "" ]; then 12 | echo "ERROR: Missing arguments." 13 | exit 14 | fi 15 | 16 | if [ ! -f $PRETRAIN ] && [ "$PRETRAIN" != "random" ]; then 17 | echo "ERROR: PRETRAIN should be a file or a string \"random\", got: $PRETRAIN" 18 | exit 19 | fi 20 | 21 | mkdir -p $WORK_DIR/logs 22 | echo "Testing pretrain: $PRETRAIN" 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 23 | 24 | bash tools/dist_extract.sh $CFG $GPUS $WORK_DIR --pretrained $PRETRAIN 25 | 26 | bash benchmarks/svm_tools/eval_svm_full.sh $WORK_DIR "$FEAT_LIST" 27 | 28 | bash benchmarks/svm_tools/eval_svm_lowshot.sh $WORK_DIR "$FEAT_LIST" 29 | -------------------------------------------------------------------------------- /benchmarks/dist_train_linear.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | set -x 5 | 6 | CFG=$1 # use cfgs under "configs/benchmarks/linear_classification/" 7 | PRETRAIN=$2 8 | PY_ARGS=${@:3} # --resume_from --deterministic 9 | GPUS=8 # When changing GPUS, please also change imgs_per_gpu in the config file accordingly to ensure the total batch size is 256. 10 | PORT=${PORT:-29500} 11 | 12 | if [ "$CFG" == "" ] || [ "$PRETRAIN" == "" ]; then 13 | echo "ERROR: Missing arguments." 14 | exit 15 | fi 16 | 17 | WORK_DIR="$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/$(echo $PRETRAIN | rev | cut -d/ -f 1 | rev)" 18 | 19 | # train 20 | python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 21 | tools/train.py \ 22 | $CFG \ 23 | --pretrained $PRETRAIN \ 24 | --work_dir $WORK_DIR --seed 0 --launcher="pytorch" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /benchmarks/dist_train_semi.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | set -x 5 | 6 | CFG=$1 # use cfgs under "configs/benchmarks/semi_classification/imagenet_*percent/" 7 | PRETRAIN=$2 8 | PY_ARGS=${@:3} 9 | GPUS=4 # in the standard setting, GPUS=4 10 | PORT=${PORT:-29500} 11 | 12 | if [ "$CFG" == "" ] || [ "$PRETRAIN" == "" ]; then 13 | echo "ERROR: Missing arguments." 14 | exit 15 | fi 16 | 17 | WORK_DIR="$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/$(echo $PRETRAIN | rev | cut -d/ -f 1 | rev)" 18 | 19 | # train 20 | python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 21 | tools/train.py \ 22 | $CFG \ 23 | --pretrained $PRETRAIN \ 24 | --work_dir $WORK_DIR --seed 0 --launcher="pytorch" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /benchmarks/extract_info/voc07.py: -------------------------------------------------------------------------------- 1 | data_source_cfg = dict(type='ImageList', memcached=False, mclient_path=None) 2 | data_root = "data/VOCdevkit/VOC2007/JPEGImages" 3 | data_all_list = "data/VOCdevkit/VOC2007/Lists/trainvaltest.txt" 4 | split_at = [5011] 5 | split_name = ['voc07_trainval', 'voc07_test'] 6 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 7 | 8 | data = dict( 9 | imgs_per_gpu=32, 10 | workers_per_gpu=2, 11 | extract=dict( 12 | type="ExtractDataset", 13 | data_source=dict( 14 | list_file=data_all_list, root=data_root, **data_source_cfg), 15 | pipeline=[ 16 | dict(type='Resize', size=256), 17 | dict(type='Resize', size=(224, 224)), 18 | dict(type='ToTensor'), 19 | dict(type='Normalize', **img_norm_cfg), 20 | ])) 21 | -------------------------------------------------------------------------------- /benchmarks/srun_test_svm_epoch.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -e 3 | set -x 4 | 5 | PARTITION=$1 6 | CFG=$2 7 | EPOCH=$3 8 | FEAT_LIST=$4 # e.g.: "feat5", "feat4 feat5". If leave empty, the default is "feat5" 9 | GPUS=${5:-8} 10 | WORK_DIR=$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/ 11 | 12 | if [ ! -f $WORK_DIR/epoch_${EPOCH}.pth ]; then 13 | echo "ERROR: File not exist: $WORK_DIR/epoch_${EPOCH}.pth" 14 | exit 15 | fi 16 | 17 | mkdir -p $WORK_DIR/logs 18 | echo "Testing checkpoint: $WORK_DIR/epoch_${EPOCH}.pth" 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 19 | 20 | bash tools/srun_extract.sh $PARTITION $CFG $GPUS $WORK_DIR --checkpoint $WORK_DIR/epoch_${EPOCH}.pth 21 | 22 | srun -p $PARTITION bash benchmarks/svm_tools/eval_svm_full.sh $WORK_DIR "$FEAT_LIST" 23 | 24 | srun -p $PARTITION bash benchmarks/svm_tools/eval_svm_lowshot.sh $WORK_DIR "$FEAT_LIST" 25 | -------------------------------------------------------------------------------- /benchmarks/srun_test_svm_pretrain.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -e 3 | set -x 4 | 5 | PARTITION=$1 6 | CFG=$2 7 | PRETRAIN=$3 # pretrained model or "random" (random init) 8 | FEAT_LIST=$4 # e.g.: "feat5", "feat4 feat5". If leave empty, the default is "feat5" 9 | GPUS=${5:-8} 10 | WORK_DIR="$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/$(echo $PRETRAIN | rev | cut -d/ -f 1 | rev)" 11 | 12 | if [ ! -f $PRETRAIN ] and [ "$PRETRAIN" != "random" ]; then 13 | echo "ERROR: PRETRAIN should be a file or a string \"random\", got: $PRETRAIN" 14 | exit 15 | fi 16 | 17 | mkdir -p $WORK_DIR/logs 18 | echo "Testing pretrain: $PRETRAIN" 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 19 | 20 | bash tools/srun_extract.sh $PARTITION $CFG $GPUS $WORK_DIR --pretrained $PRETRAIN 21 | 22 | srun -p $PARTITION bash benchmarks/svm_tools/eval_svm_full.sh $WORK_DIR "$FEAT_LIST" 23 | 24 | srun -p $PARTITION bash benchmarks/svm_tools/eval_svm_lowshot.sh $WORK_DIR "$FEAT_LIST" 25 | -------------------------------------------------------------------------------- /benchmarks/srun_train_linear.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | set -x 5 | 6 | PARTITION=$1 7 | CFG=$2 8 | PRETRAIN=$3 9 | PY_ARGS=${@:4} 10 | JOB_NAME="openselfsup" 11 | GPUS=8 # When changing GPUS, please also change imgs_per_gpu in the config file accordingly to ensure the total batch size is 256. 12 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 13 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 14 | SRUN_ARGS=${SRUN_ARGS:-""} 15 | 16 | WORK_DIR="$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/$(echo $PRETRAIN | rev | cut -d/ -f 1 | rev)" 17 | 18 | # train 19 | GLOG_vmodule=MemcachedClient=-1 \ 20 | srun -p ${PARTITION} \ 21 | --job-name=${JOB_NAME} \ 22 | --gres=gpu:${GPUS_PER_NODE} \ 23 | --ntasks=${GPUS} \ 24 | --ntasks-per-node=${GPUS_PER_NODE} \ 25 | --cpus-per-task=${CPUS_PER_TASK} \ 26 | --kill-on-bad-exit=1 \ 27 | ${SRUN_ARGS} \ 28 | python -u tools/train.py \ 29 | $CFG \ 30 | --pretrained $PRETRAIN \ 31 | --work_dir $WORK_DIR --seed 0 --launcher="slurm" ${PY_ARGS} 32 | -------------------------------------------------------------------------------- /benchmarks/srun_train_semi.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | set -x 5 | 6 | PARTITION=$1 7 | CFG=$2 8 | PRETRAIN=$3 9 | PY_ARGS=${@:4} 10 | JOB_NAME="openselfsup" 11 | GPUS=4 # in the standard setting, GPUS=4 12 | GPUS_PER_NODE=${GPUS_PER_NODE:-4} 13 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 14 | SRUN_ARGS=${SRUN_ARGS:-""} 15 | 16 | WORK_DIR="$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/$(echo $PRETRAIN | rev | cut -d/ -f 1 | rev)" 17 | 18 | # train 19 | GLOG_vmodule=MemcachedClient=-1 \ 20 | srun -p ${PARTITION} \ 21 | --job-name=${JOB_NAME} \ 22 | --gres=gpu:${GPUS_PER_NODE} \ 23 | --ntasks=${GPUS} \ 24 | --ntasks-per-node=${GPUS_PER_NODE} \ 25 | --cpus-per-task=${CPUS_PER_TASK} \ 26 | --kill-on-bad-exit=1 \ 27 | ${SRUN_ARGS} \ 28 | python -u tools/train.py \ 29 | $CFG \ 30 | --pretrained $PRETRAIN \ 31 | --work_dir $WORK_DIR --seed 0 --launcher="slurm" ${PY_ARGS} 32 | -------------------------------------------------------------------------------- /benchmarks/svm_tools/eval_svm_full.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -x 3 | set -e 4 | 5 | WORK_DIR=$1 6 | FEAT_LIST=${2:-"feat5"} # "feat1 feat2 feat3 feat4 feat5" 7 | TRAIN_SVM_FLAG=true 8 | TEST_SVM_FLAG=true 9 | DATA="data/VOCdevkit/VOC2007/SVMLabels" 10 | 11 | # config svm 12 | costs="1.0,10.0,100.0" 13 | 14 | for feat in $FEAT_LIST; do 15 | echo "For feature: $feat" 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 16 | # train svm 17 | if $TRAIN_SVM_FLAG; then 18 | rm -rf $WORK_DIR/svm 19 | mkdir -p $WORK_DIR/svm/voc07_${feat} 20 | echo "training svm ..." 21 | python benchmarks/svm_tools/train_svm_kfold_parallel.py \ 22 | --data_file $WORK_DIR/features/voc07_trainval_${feat}.npy \ 23 | --targets_data_file $DATA/train_labels.npy \ 24 | --costs_list $costs \ 25 | --output_path $WORK_DIR/svm/voc07_${feat} 26 | fi 27 | 28 | # test svm 29 | if $TEST_SVM_FLAG; then 30 | echo "testing svm ..." 31 | python benchmarks/svm_tools/test_svm.py \ 32 | --data_file $WORK_DIR/features/voc07_test_${feat}.npy \ 33 | --json_targets $DATA/test_targets.json \ 34 | --targets_data_file $DATA/test_labels.npy \ 35 | --costs_list $costs \ 36 | --generate_json 1 \ 37 | --output_path $WORK_DIR/svm/voc07_${feat} 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 38 | fi 39 | 40 | done 41 | -------------------------------------------------------------------------------- /benchmarks/svm_tools/eval_svm_lowshot.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -x 3 | set -e 4 | 5 | WORK_DIR=$1 6 | MODE="full" 7 | FEAT_LIST=${2:-"feat5"} # "feat1 feat2 feat3 feat4 feat5" 8 | TRAIN_SVM_LOWSHOT_FLAG=true 9 | TEST_SVM_LOWSHOT_FLAG=true 10 | AGGREGATE_FLAG=true 11 | DATA="data/VOCdevkit/VOC2007/SVMLabels" 12 | 13 | # config svm 14 | costs="1.0,10.0,100.0" 15 | if [ "$MODE" == "fast" ]; then 16 | shots="96" 17 | else 18 | shots="1 2 4 8 16 32 64 96" 19 | fi 20 | 21 | for feat in $FEAT_LIST; do 22 | echo "For feature: $feat" 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 23 | # train lowshot svm 24 | if $TRAIN_SVM_LOWSHOT_FLAG; then 25 | rm -rf $WORK_DIR/svm_lowshot 26 | mkdir -p $WORK_DIR/svm_lowshot/voc07_${feat} 27 | echo "training svm low-shot ..." 28 | for s in {1..5}; do 29 | for k in $shots; do 30 | echo -e "\ts${s} k${k}" 31 | python benchmarks/svm_tools/train_svm_low_shot.py \ 32 | --data_file $WORK_DIR/features/voc07_trainval_${feat}.npy \ 33 | --targets_data_file $DATA/low_shot/labels/train_targets_sample${s}_k${k}.npy \ 34 | --costs_list $costs \ 35 | --output_path $WORK_DIR/svm_lowshot/voc07_${feat} 36 | done 37 | done 38 | fi 39 | 40 | # test lowshot svm 41 | if $TEST_SVM_LOWSHOT_FLAG; then 42 | echo "testing svm low-shot ..." 43 | python benchmarks/svm_tools/test_svm_low_shot.py \ 44 | --data_file $WORK_DIR/features/voc07_test_${feat}.npy \ 45 | --targets_data_file $DATA/test_labels.npy \ 46 | --json_targets $DATA/test_targets.json \ 47 | --generate_json 1 \ 48 | --costs_list $costs \ 49 | --output_path $WORK_DIR/svm_lowshot/voc07_${feat} \ 50 | --k_values "${shots// /,}" \ 51 | --sample_inds "0,1,2,3,4" \ 52 | --dataset "voc" 53 | fi 54 | 55 | # aggregate testing results 56 | if $AGGREGATE_FLAG; then 57 | echo "aggregating svm low-shot ..." 58 | python benchmarks/svm_tools/aggregate_low_shot_svm_stats.py \ 59 | --output_path $WORK_DIR/svm_lowshot/voc07_${feat} \ 60 | --k_values "${shots// /,}" \ 61 | --sample_inds "0,1,2,3,4" 2>&1 | tee -a $WORK_DIR/logs/eval_svm.log 62 | fi 63 | 64 | done 65 | -------------------------------------------------------------------------------- /configs/base.py: -------------------------------------------------------------------------------- 1 | train_cfg = {} 2 | test_cfg = {} 3 | optimizer_config = dict() # grad_clip, coalesce, bucket_size_mb 4 | # yapf:disable 5 | log_config = dict( 6 | interval=50, 7 | hooks=[ 8 | dict(type='TextLoggerHook'), 9 | dict(type='TensorboardLoggerHook') 10 | ]) 11 | # yapf:enable 12 | # runtime settings 13 | dist_params = dict(backend='nccl') 14 | cudnn_benchmark = True 15 | log_level = 'INFO' 16 | load_from = None 17 | resume_from = None 18 | workflow = [('train', 1)] 19 | prefetch = False 20 | -------------------------------------------------------------------------------- /configs/benchmarks/linear_classification/imagenet/r50_last.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=False, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=3, 11 | out_indices=[4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='BN'), 13 | frozen_stages=4), 14 | head=dict( 15 | type='ClsHead', with_avg_pool=True, in_channels=2048, 16 | num_classes=1000)) 17 | # dataset settings 18 | data_source_cfg = dict( 19 | type='ImageNet', 20 | memcached=True, 21 | mclient_path='/mnt/lustre/share/memcached_client') 22 | data_train_list = 'data/imagenet/meta/train_labeled.txt' 23 | data_train_root = 'data/imagenet/train' 24 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 25 | data_test_root = 'data/imagenet/val' 26 | dataset_type = 'ClassificationDataset' 27 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 28 | train_pipeline = [ 29 | dict(type='RandomResizedCrop', size=224), 30 | dict(type='RandomHorizontalFlip'), 31 | ] 32 | test_pipeline = [ 33 | dict(type='Resize', size=256), 34 | dict(type='CenterCrop', size=224), 35 | ] 36 | # prefetch 37 | prefetch = False 38 | if not prefetch: 39 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 40 | test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 41 | data = dict( 42 | imgs_per_gpu=32, # total 32*8=256, 8GPU linear cls 43 | workers_per_gpu=5, 44 | train=dict( 45 | type=dataset_type, 46 | data_source=dict( 47 | list_file=data_train_list, root=data_train_root, 48 | **data_source_cfg), 49 | pipeline=train_pipeline, 50 | prefetch=prefetch), 51 | val=dict( 52 | type=dataset_type, 53 | data_source=dict( 54 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 55 | pipeline=test_pipeline, 56 | prefetch=prefetch)) 57 | # additional hooks 58 | custom_hooks = [ 59 | dict( 60 | type='ValidateHook', 61 | dataset=data['val'], 62 | initial=True, 63 | interval=1, 64 | imgs_per_gpu=128, 65 | workers_per_gpu=4, 66 | prefetch=prefetch, 67 | img_norm_cfg=img_norm_cfg, 68 | eval_param=dict(topk=(1, 5))) 69 | ] 70 | # optimizer 71 | optimizer = dict(type='SGD', lr=30., momentum=0.9, weight_decay=0.) 72 | # learning policy 73 | lr_config = dict(policy='step', step=[60, 80]) 74 | checkpoint_config = dict(interval=10) 75 | # runtime settings 76 | total_epochs = 100 77 | -------------------------------------------------------------------------------- /configs/benchmarks/linear_classification/imagenet/r50_last_sobel.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=True, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=2, 11 | out_indices=[4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='BN'), 13 | frozen_stages=4), 14 | head=dict( 15 | type='ClsHead', with_avg_pool=True, in_channels=2048, 16 | num_classes=1000)) 17 | # dataset settings 18 | data_source_cfg = dict( 19 | type='ImageNet', 20 | memcached=True, 21 | mclient_path='/mnt/lustre/share/memcached_client') 22 | data_train_list = 'data/imagenet/meta/train_labeled.txt' 23 | data_train_root = 'data/imagenet/train' 24 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 25 | data_test_root = 'data/imagenet/val' 26 | dataset_type = 'ClassificationDataset' 27 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 28 | train_pipeline = [ 29 | dict(type='RandomResizedCrop', size=224), 30 | dict(type='RandomHorizontalFlip'), 31 | ] 32 | test_pipeline = [ 33 | dict(type='Resize', size=256), 34 | dict(type='CenterCrop', size=224), 35 | ] 36 | # prefetch 37 | prefetch = False 38 | if not prefetch: 39 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 40 | test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 41 | data = dict( 42 | imgs_per_gpu=32, # total 32*8=256, 8GPU linear cls 43 | workers_per_gpu=5, 44 | train=dict( 45 | type=dataset_type, 46 | data_source=dict( 47 | list_file=data_train_list, root=data_train_root, 48 | **data_source_cfg), 49 | pipeline=train_pipeline, 50 | prefetch=prefetch), 51 | val=dict( 52 | type=dataset_type, 53 | data_source=dict( 54 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 55 | pipeline=test_pipeline, 56 | prefetch=prefetch)) 57 | # additional hooks 58 | custom_hooks = [ 59 | dict( 60 | type='ValidateHook', 61 | dataset=data['val'], 62 | initial=True, 63 | interval=1, 64 | imgs_per_gpu=128, 65 | workers_per_gpu=4, 66 | prefetch=prefetch, 67 | img_norm_cfg=img_norm_cfg, 68 | eval_param=dict(topk=(1, 5))) 69 | ] 70 | # optimizer 71 | optimizer = dict(type='SGD', lr=30., momentum=0.9, weight_decay=0.) 72 | # learning policy 73 | lr_config = dict(policy='step', step=[60, 80]) 74 | checkpoint_config = dict(interval=10) 75 | # runtime settings 76 | total_epochs = 100 77 | -------------------------------------------------------------------------------- /configs/benchmarks/linear_classification/imagenet/r50_multihead.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=False, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=3, 11 | out_indices=[0, 1, 2, 3, 4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='BN'), 13 | frozen_stages=4), 14 | head=dict( 15 | type='MultiClsHead', 16 | pool_type='specified', 17 | in_indices=[0, 1, 2, 3, 4], 18 | with_last_layer_unpool=False, 19 | backbone='resnet50', 20 | norm_cfg=dict(type='SyncBN', momentum=0.1, affine=False), 21 | num_classes=1000)) 22 | # dataset settings 23 | data_source_cfg = dict( 24 | type='ImageNet', 25 | memcached=True, 26 | mclient_path='/mnt/lustre/share/memcached_client') 27 | data_train_list = 'data/imagenet/meta/train_labeled.txt' 28 | data_train_root = 'data/imagenet/train' 29 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 30 | data_test_root = 'data/imagenet/val' 31 | dataset_type = 'ClassificationDataset' 32 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 33 | train_pipeline = [ 34 | dict(type='RandomResizedCrop', size=224), 35 | dict(type='RandomHorizontalFlip'), 36 | dict( 37 | type='ColorJitter', 38 | brightness=0.4, 39 | contrast=0.4, 40 | saturation=0.4, 41 | hue=0.), 42 | dict(type='ToTensor'), 43 | dict(type='Lighting'), 44 | dict(type='Normalize', **img_norm_cfg), 45 | ] 46 | test_pipeline = [ 47 | dict(type='Resize', size=256), 48 | dict(type='CenterCrop', size=224), 49 | dict(type='ToTensor'), 50 | dict(type='Normalize', **img_norm_cfg), 51 | ] 52 | data = dict( 53 | imgs_per_gpu=32, # total 32x8=256 54 | workers_per_gpu=5, 55 | train=dict( 56 | type=dataset_type, 57 | data_source=dict( 58 | list_file=data_train_list, root=data_train_root, 59 | **data_source_cfg), 60 | pipeline=train_pipeline), 61 | val=dict( 62 | type=dataset_type, 63 | data_source=dict( 64 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 65 | pipeline=test_pipeline)) 66 | # additional hooks 67 | custom_hooks = [ 68 | dict( 69 | type='ValidateHook', 70 | dataset=data['val'], 71 | initial=True, 72 | interval=10, 73 | imgs_per_gpu=128, 74 | workers_per_gpu=4, 75 | eval_param=dict(topk=(1, ))) 76 | ] 77 | # optimizer 78 | optimizer = dict( 79 | type='SGD', 80 | lr=0.01, 81 | momentum=0.9, 82 | weight_decay=0.0001, 83 | paramwise_options=dict(norm_decay_mult=0.), 84 | nesterov=True) 85 | # learning policy 86 | lr_config = dict(policy='step', step=[30, 60, 90]) 87 | checkpoint_config = dict(interval=10) 88 | # runtime settings 89 | total_epochs = 90 90 | -------------------------------------------------------------------------------- /configs/benchmarks/linear_classification/imagenet/r50_multihead_sobel.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=True, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=2, 11 | out_indices=[0, 1, 2, 3, 4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='BN'), 13 | frozen_stages=4), 14 | head=dict( 15 | type='MultiClsHead', 16 | pool_type='specified', 17 | in_indices=[0, 1, 2, 3, 4], 18 | with_last_layer_unpool=False, 19 | backbone='resnet50', 20 | norm_cfg=dict(type='SyncBN', momentum=0.1, affine=False), 21 | num_classes=1000)) 22 | # dataset settings 23 | data_source_cfg = dict( 24 | type='ImageNet', 25 | memcached=True, 26 | mclient_path='/mnt/lustre/share/memcached_client') 27 | data_train_list = 'data/imagenet/meta/train_labeled.txt' 28 | data_train_root = 'data/imagenet/train' 29 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 30 | data_test_root = 'data/imagenet/val' 31 | dataset_type = 'ClassificationDataset' 32 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 33 | train_pipeline = [ 34 | dict(type='RandomResizedCrop', size=224), 35 | dict(type='RandomHorizontalFlip'), 36 | dict( 37 | type='ColorJitter', 38 | brightness=0.4, 39 | contrast=0.4, 40 | saturation=0.4, 41 | hue=0.), 42 | dict(type='ToTensor'), 43 | dict(type='Lighting'), 44 | dict(type='Normalize', **img_norm_cfg), 45 | ] 46 | test_pipeline = [ 47 | dict(type='Resize', size=256), 48 | dict(type='CenterCrop', size=224), 49 | dict(type='ToTensor'), 50 | dict(type='Normalize', **img_norm_cfg), 51 | ] 52 | data = dict( 53 | imgs_per_gpu=32, # total 32x8=256 54 | workers_per_gpu=5, 55 | train=dict( 56 | type=dataset_type, 57 | data_source=dict( 58 | list_file=data_train_list, root=data_train_root, 59 | **data_source_cfg), 60 | pipeline=train_pipeline), 61 | val=dict( 62 | type=dataset_type, 63 | data_source=dict( 64 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 65 | pipeline=test_pipeline)) 66 | # additional hooks 67 | custom_hooks = [ 68 | dict( 69 | type='ValidateHook', 70 | dataset=data['val'], 71 | initial=True, 72 | interval=10, 73 | imgs_per_gpu=128, 74 | workers_per_gpu=4, 75 | eval_param=dict(topk=(1, ))) 76 | ] 77 | # optimizer 78 | optimizer = dict( 79 | type='SGD', 80 | lr=0.01, 81 | momentum=0.9, 82 | weight_decay=0.0001, 83 | paramwise_options=dict(norm_decay_mult=0.), 84 | nesterov=True) 85 | # learning policy 86 | lr_config = dict(policy='step', step=[30, 60, 90]) 87 | checkpoint_config = dict(interval=10) 88 | # runtime settings 89 | total_epochs = 90 90 | -------------------------------------------------------------------------------- /configs/benchmarks/linear_classification/places205/r50_multihead.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=False, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=3, 11 | out_indices=[0, 1, 2, 3, 4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='BN'), 13 | frozen_stages=4), 14 | head=dict( 15 | type='MultiClsHead', 16 | pool_type='specified', 17 | in_indices=[0, 1, 2, 3, 4], 18 | with_last_layer_unpool=False, 19 | backbone='resnet50', 20 | norm_cfg=dict(type='SyncBN', momentum=0.1, affine=False), 21 | num_classes=205)) 22 | # dataset settings 23 | data_source_cfg = dict( 24 | type='Places205', 25 | memcached=True, 26 | mclient_path='/mnt/lustre/share/memcached_client') 27 | data_train_list = 'data/places205/meta/train_labeled.txt' 28 | data_train_root = 'data/places205/train' 29 | data_test_list = 'data/places205/meta/val_labeled.txt' 30 | data_test_root = 'data/places205/val' 31 | dataset_type = 'ClassificationDataset' 32 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 33 | train_pipeline = [ 34 | dict(type='Resize', size=256), 35 | dict(type='CenterCrop', size=256), 36 | dict(type='RandomCrop', size=224), 37 | dict(type='RandomHorizontalFlip'), 38 | ] 39 | test_pipeline = [ 40 | dict(type='Resize', size=256), 41 | dict(type='CenterCrop', size=224), 42 | ] 43 | # prefetch 44 | prefetch = False 45 | if not prefetch: 46 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 47 | test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 48 | data = dict( 49 | imgs_per_gpu=32, # total 32x8=256 50 | workers_per_gpu=4, 51 | train=dict( 52 | type=dataset_type, 53 | data_source=dict( 54 | list_file=data_train_list, root=data_train_root, 55 | **data_source_cfg), 56 | pipeline=train_pipeline, 57 | prefetch=prefetch), 58 | val=dict( 59 | type=dataset_type, 60 | data_source=dict( 61 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 62 | pipeline=test_pipeline, 63 | prefetch=prefetch)) 64 | # additional hooks 65 | custom_hooks = [ 66 | dict( 67 | type='ValidateHook', 68 | dataset=data['val'], 69 | initial=True, 70 | interval=10, 71 | imgs_per_gpu=32, 72 | workers_per_gpu=4, 73 | prefetch=prefetch, 74 | img_norm_cfg=img_norm_cfg, 75 | eval_param=dict(topk=(1, ))) 76 | ] 77 | # optimizer 78 | optimizer = dict( 79 | type='SGD', 80 | lr=0.01, 81 | momentum=0.9, 82 | weight_decay=0.0001, 83 | paramwise_options=dict(norm_decay_mult=0.), 84 | nesterov=True) 85 | # learning policy 86 | lr_config = dict(policy='step', step=[7, 14, 21]) 87 | checkpoint_config = dict(interval=10) 88 | # runtime settings 89 | total_epochs = 28 90 | -------------------------------------------------------------------------------- /configs/benchmarks/linear_classification/places205/r50_multihead_sobel.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=True, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=2, 11 | out_indices=[0, 1, 2, 3, 4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='BN'), 13 | frozen_stages=4), 14 | head=dict( 15 | type='MultiClsHead', 16 | pool_type='specified', 17 | in_indices=[0, 1, 2, 3, 4], 18 | with_last_layer_unpool=False, 19 | backbone='resnet50', 20 | norm_cfg=dict(type='SyncBN', momentum=0.1, affine=False), 21 | num_classes=205)) 22 | # dataset settings 23 | data_source_cfg = dict( 24 | type='Places205', 25 | memcached=True, 26 | mclient_path='/mnt/lustre/share/memcached_client') 27 | data_train_list = 'data/places205/meta/train_labeled.txt' 28 | data_train_root = 'data/places205/train' 29 | data_test_list = 'data/places205/meta/val_labeled.txt' 30 | data_test_root = 'data/places205/val' 31 | dataset_type = 'ClassificationDataset' 32 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 33 | train_pipeline = [ 34 | dict(type='Resize', size=256), 35 | dict(type='CenterCrop', size=256), 36 | dict(type='RandomCrop', size=224), 37 | dict(type='RandomHorizontalFlip'), 38 | ] 39 | test_pipeline = [ 40 | dict(type='Resize', size=256), 41 | dict(type='CenterCrop', size=224), 42 | ] 43 | # prefetch 44 | prefetch = False 45 | if not prefetch: 46 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 47 | test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 48 | data = dict( 49 | imgs_per_gpu=32, # total 32x8=256 50 | workers_per_gpu=4, 51 | train=dict( 52 | type=dataset_type, 53 | data_source=dict( 54 | list_file=data_train_list, root=data_train_root, 55 | **data_source_cfg), 56 | pipeline=train_pipeline, 57 | prefetch=prefetch), 58 | val=dict( 59 | type=dataset_type, 60 | data_source=dict( 61 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 62 | pipeline=test_pipeline, 63 | prefetch=prefetch)) 64 | # additional hooks 65 | custom_hooks = [ 66 | dict( 67 | type='ValidateHook', 68 | dataset=data['val'], 69 | initial=True, 70 | interval=10, 71 | imgs_per_gpu=32, 72 | workers_per_gpu=4, 73 | prefetch=prefetch, 74 | img_norm_cfg=img_norm_cfg, 75 | eval_param=dict(topk=(1, ))) 76 | ] 77 | # optimizer 78 | optimizer = dict( 79 | type='SGD', 80 | lr=0.01, 81 | momentum=0.9, 82 | weight_decay=0.0001, 83 | paramwise_options=dict(norm_decay_mult=0.), 84 | nesterov=True) 85 | # learning policy 86 | lr_config = dict(policy='step', step=[7, 14, 21]) 87 | checkpoint_config = dict(interval=10) 88 | # runtime settings 89 | total_epochs = 28 90 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/base.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | out_indices=[4], # 0: conv-1, x: stage-x 10 | norm_cfg=dict(type='SyncBN')), 11 | head=dict( 12 | type='ClsHead', with_avg_pool=True, in_channels=2048, 13 | num_classes=1000)) 14 | # dataset settings 15 | data_source_cfg = dict( 16 | type='ImageNet', 17 | memcached=True, 18 | mclient_path='/mnt/lustre/share/memcached_client') 19 | data_train_list = 'data/imagenet/meta/train_labeled_10percent.txt' 20 | data_train_root = 'data/imagenet/train' 21 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 22 | data_test_root = 'data/imagenet/val' 23 | dataset_type = 'ClassificationDataset' 24 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 25 | train_pipeline = [ 26 | dict(type='RandomResizedCrop', size=224), 27 | dict(type='RandomHorizontalFlip'), 28 | dict(type='ToTensor'), 29 | dict(type='Normalize', **img_norm_cfg), 30 | ] 31 | test_pipeline = [ 32 | dict(type='Resize', size=256), 33 | dict(type='CenterCrop', size=224), 34 | dict(type='ToTensor'), 35 | dict(type='Normalize', **img_norm_cfg), 36 | ] 37 | data = dict( 38 | imgs_per_gpu=64, # total 256 39 | workers_per_gpu=2, 40 | train=dict( 41 | type=dataset_type, 42 | data_source=dict( 43 | list_file=data_train_list, root=data_train_root, 44 | **data_source_cfg), 45 | pipeline=train_pipeline), 46 | val=dict( 47 | type=dataset_type, 48 | data_source=dict( 49 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 50 | pipeline=test_pipeline)) 51 | # additional hooks 52 | custom_hooks = [ 53 | dict( 54 | type='ValidateHook', 55 | dataset=data['val'], 56 | initial=False, 57 | interval=20, 58 | imgs_per_gpu=32, 59 | workers_per_gpu=2, 60 | eval_param=dict(topk=(1, 5))) 61 | ] 62 | # learning policy 63 | lr_config = dict(policy='step', step=[12, 16], gamma=0.2) 64 | checkpoint_config = dict(interval=20) 65 | # runtime settings 66 | total_epochs = 20 67 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_001_head1.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_001_head10.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=10)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_001_head100.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=100)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_01_head1.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_01_head10.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=10)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_01_head100.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=100)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_01_head1_sobel.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=True, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=2, 11 | out_indices=[4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='SyncBN')), 13 | head=dict( 14 | type='ClsHead', with_avg_pool=True, in_channels=2048, 15 | num_classes=1000)) 16 | # dataset settings 17 | data_source_cfg = dict( 18 | type='ImageNet', 19 | memcached=True, 20 | mclient_path='/mnt/lustre/share/memcached_client') 21 | data_train_list = 'data/imagenet/meta/train_labeled_10percent.txt' 22 | data_train_root = 'data/imagenet/train' 23 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 24 | data_test_root = 'data/imagenet/val' 25 | dataset_type = 'ClassificationDataset' 26 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 27 | train_pipeline = [ 28 | dict(type='RandomResizedCrop', size=224), 29 | dict(type='RandomHorizontalFlip'), 30 | dict(type='ToTensor'), 31 | dict(type='Normalize', **img_norm_cfg), 32 | ] 33 | test_pipeline = [ 34 | dict(type='Resize', size=256), 35 | dict(type='CenterCrop', size=224), 36 | dict(type='ToTensor'), 37 | dict(type='Normalize', **img_norm_cfg), 38 | ] 39 | data = dict( 40 | imgs_per_gpu=64, # total 256 41 | workers_per_gpu=2, 42 | train=dict( 43 | type=dataset_type, 44 | data_source=dict( 45 | list_file=data_train_list, root=data_train_root, 46 | **data_source_cfg), 47 | pipeline=train_pipeline), 48 | val=dict( 49 | type=dataset_type, 50 | data_source=dict( 51 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 52 | pipeline=test_pipeline)) 53 | # additional hooks 54 | custom_hooks = [ 55 | dict( 56 | type='ValidateHook', 57 | dataset=data['val'], 58 | initial=False, 59 | interval=20, 60 | imgs_per_gpu=32, 61 | workers_per_gpu=2, 62 | eval_param=dict(topk=(1, 5))) 63 | ] 64 | # optimizer 65 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001, 66 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 67 | # learning policy 68 | lr_config = dict(policy='step', step=[12, 16], gamma=0.2) 69 | checkpoint_config = dict(interval=20) 70 | # runtime settings 71 | total_epochs = 20 72 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_1_head1.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_1_head10.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=10)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_10percent/r50_lr0_1_head100.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001, 4 | paramwise_options={'\Ahead.': dict(lr_mult=100)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/base.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | out_indices=[4], # 0: conv-1, x: stage-x 10 | norm_cfg=dict(type='SyncBN')), 11 | head=dict( 12 | type='ClsHead', with_avg_pool=True, in_channels=2048, 13 | num_classes=1000)) 14 | # dataset settings 15 | data_source_cfg = dict( 16 | type='ImageNet', 17 | memcached=True, 18 | mclient_path='/mnt/lustre/share/memcached_client') 19 | data_train_list = 'data/imagenet/meta/train_labeled_1percent.txt' 20 | data_train_root = 'data/imagenet/train' 21 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 22 | data_test_root = 'data/imagenet/val' 23 | dataset_type = 'ClassificationDataset' 24 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 25 | train_pipeline = [ 26 | dict(type='RandomResizedCrop', size=224), 27 | dict(type='RandomHorizontalFlip'), 28 | dict(type='ToTensor'), 29 | dict(type='Normalize', **img_norm_cfg), 30 | ] 31 | test_pipeline = [ 32 | dict(type='Resize', size=256), 33 | dict(type='CenterCrop', size=224), 34 | dict(type='ToTensor'), 35 | dict(type='Normalize', **img_norm_cfg), 36 | ] 37 | data = dict( 38 | imgs_per_gpu=64, # total 256 39 | workers_per_gpu=2, 40 | train=dict( 41 | type=dataset_type, 42 | data_source=dict( 43 | list_file=data_train_list, root=data_train_root, 44 | **data_source_cfg), 45 | pipeline=train_pipeline), 46 | val=dict( 47 | type=dataset_type, 48 | data_source=dict( 49 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 50 | pipeline=test_pipeline)) 51 | # additional hooks 52 | custom_hooks = [ 53 | dict( 54 | type='ValidateHook', 55 | dataset=data['val'], 56 | initial=False, 57 | interval=20, 58 | imgs_per_gpu=32, 59 | workers_per_gpu=2, 60 | eval_param=dict(topk=(1, 5))) 61 | ] 62 | # learning policy 63 | lr_config = dict(policy='step', step=[12, 16], gamma=0.2) 64 | checkpoint_config = dict(interval=20) 65 | log_config = dict( 66 | interval=10, 67 | hooks=[ 68 | dict(type='TextLoggerHook'), 69 | dict(type='TensorboardLoggerHook') 70 | ]) 71 | # runtime settings 72 | total_epochs = 20 73 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_001_head1.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_001_head10.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=10)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_001_head100.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=100)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_01_head1.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_01_head10.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=10)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_01_head100.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=100)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_01_head1_sobel.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | with_sobel=True, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=2, 11 | out_indices=[4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='SyncBN')), 13 | head=dict( 14 | type='ClsHead', with_avg_pool=True, in_channels=2048, 15 | num_classes=1000)) 16 | # dataset settings 17 | data_source_cfg = dict( 18 | type='ImageNet', 19 | memcached=True, 20 | mclient_path='/mnt/lustre/share/memcached_client') 21 | data_train_list = 'data/imagenet/meta/train_labeled_1percent.txt' 22 | data_train_root = 'data/imagenet/train' 23 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 24 | data_test_root = 'data/imagenet/val' 25 | dataset_type = 'ClassificationDataset' 26 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 27 | train_pipeline = [ 28 | dict(type='RandomResizedCrop', size=224), 29 | dict(type='RandomHorizontalFlip'), 30 | dict(type='ToTensor'), 31 | dict(type='Normalize', **img_norm_cfg), 32 | ] 33 | test_pipeline = [ 34 | dict(type='Resize', size=256), 35 | dict(type='CenterCrop', size=224), 36 | dict(type='ToTensor'), 37 | dict(type='Normalize', **img_norm_cfg), 38 | ] 39 | data = dict( 40 | imgs_per_gpu=64, # total 256 41 | workers_per_gpu=2, 42 | train=dict( 43 | type=dataset_type, 44 | data_source=dict( 45 | list_file=data_train_list, root=data_train_root, 46 | **data_source_cfg), 47 | pipeline=train_pipeline), 48 | val=dict( 49 | type=dataset_type, 50 | data_source=dict( 51 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 52 | pipeline=test_pipeline)) 53 | # additional hooks 54 | custom_hooks = [ 55 | dict( 56 | type='ValidateHook', 57 | dataset=data['val'], 58 | initial=False, 59 | interval=20, 60 | imgs_per_gpu=32, 61 | workers_per_gpu=2, 62 | eval_param=dict(topk=(1, 5))) 63 | ] 64 | # optimizer 65 | optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005, 66 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 67 | # learning policy 68 | lr_config = dict(policy='step', step=[12, 16], gamma=0.2) 69 | checkpoint_config = dict(interval=20) 70 | log_config = dict( 71 | interval=10, 72 | hooks=[ 73 | dict(type='TextLoggerHook'), 74 | dict(type='TensorboardLoggerHook') 75 | ]) 76 | # runtime settings 77 | total_epochs = 20 78 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_1_head1.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=1)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_1_head10.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=10)}) 5 | -------------------------------------------------------------------------------- /configs/benchmarks/semi_classification/imagenet_1percent/r50_lr0_1_head100.py: -------------------------------------------------------------------------------- 1 | _base_ = 'base.py' 2 | # optimizer 3 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0005, 4 | paramwise_options={'\Ahead.': dict(lr_mult=100)}) 5 | -------------------------------------------------------------------------------- /configs/classification/cifar10/r50.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | out_indices=[4], # 4: stage-4 10 | norm_cfg=dict(type='BN')), 11 | head=dict( 12 | type='ClsHead', with_avg_pool=True, in_channels=2048, num_classes=10)) 13 | # dataset settings 14 | data_source_cfg = dict(type='Cifar10', root='data/cifar/') 15 | dataset_type = 'ClassificationDataset' 16 | img_norm_cfg = dict(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.201]) 17 | train_pipeline = [ 18 | dict(type='RandomCrop', size=32, padding=4), 19 | dict(type='RandomHorizontalFlip'), 20 | dict(type='ToTensor'), 21 | dict(type='Normalize', **img_norm_cfg), 22 | ] 23 | test_pipeline = [ 24 | dict(type='ToTensor'), 25 | dict(type='Normalize', **img_norm_cfg), 26 | ] 27 | data = dict( 28 | imgs_per_gpu=128, 29 | workers_per_gpu=2, 30 | train=dict( 31 | type=dataset_type, 32 | data_source=dict(split='train', **data_source_cfg), 33 | pipeline=train_pipeline), 34 | val=dict( 35 | type=dataset_type, 36 | data_source=dict(split='test', **data_source_cfg), 37 | pipeline=test_pipeline), 38 | test=dict( 39 | type=dataset_type, 40 | data_source=dict(split='test', **data_source_cfg), 41 | pipeline=test_pipeline)) 42 | # additional hooks 43 | custom_hooks = [ 44 | dict( 45 | type='ValidateHook', 46 | dataset=data['val'], 47 | initial=True, 48 | interval=10, 49 | imgs_per_gpu=128, 50 | workers_per_gpu=8, 51 | eval_param=dict(topk=(1, 5))) 52 | ] 53 | # optimizer 54 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0005) 55 | # learning policy 56 | lr_config = dict(policy='step', step=[150, 250]) 57 | checkpoint_config = dict(interval=50) 58 | # runtime settings 59 | total_epochs = 350 60 | -------------------------------------------------------------------------------- /configs/classification/imagenet/r50.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='Classification', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | out_indices=[4], # 0: conv-1, x: stage-x 10 | norm_cfg=dict(type='SyncBN')), 11 | head=dict( 12 | type='ClsHead', with_avg_pool=True, in_channels=2048, 13 | num_classes=1000)) 14 | # dataset settings 15 | data_source_cfg = dict( 16 | type='ImageNet', 17 | memcached=True, 18 | mclient_path='/mnt/lustre/share/memcached_client') 19 | data_train_list = 'data/imagenet/meta/train_labeled.txt' 20 | data_train_root = 'data/imagenet/train' 21 | data_test_list = 'data/imagenet/meta/val_labeled.txt' 22 | data_test_root = 'data/imagenet/val' 23 | dataset_type = 'ClassificationDataset' 24 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 25 | train_pipeline = [ 26 | dict(type='RandomResizedCrop', size=224), 27 | dict(type='RandomHorizontalFlip'), 28 | dict(type='ToTensor'), 29 | dict(type='Normalize', **img_norm_cfg), 30 | ] 31 | test_pipeline = [ 32 | dict(type='Resize', size=256), 33 | dict(type='CenterCrop', size=224), 34 | dict(type='ToTensor'), 35 | dict(type='Normalize', **img_norm_cfg), 36 | ] 37 | data = dict( 38 | imgs_per_gpu=32, # total 256 39 | workers_per_gpu=2, 40 | train=dict( 41 | type=dataset_type, 42 | data_source=dict( 43 | list_file=data_train_list, root=data_train_root, 44 | **data_source_cfg), 45 | pipeline=train_pipeline), 46 | val=dict( 47 | type=dataset_type, 48 | data_source=dict( 49 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 50 | pipeline=test_pipeline)) 51 | # additional hooks 52 | custom_hooks = [ 53 | dict( 54 | type='ValidateHook', 55 | dataset=data['val'], 56 | initial=True, 57 | interval=10, 58 | imgs_per_gpu=32, 59 | workers_per_gpu=2, 60 | eval_param=dict(topk=(1, 5))) 61 | ] 62 | # optimizer 63 | optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001) 64 | # learning policy 65 | lr_config = dict(policy='step', step=[30, 60, 90]) 66 | checkpoint_config = dict(interval=10) 67 | # runtime settings 68 | total_epochs = 90 69 | -------------------------------------------------------------------------------- /configs/selfsup/byol/r50_bs256_ep200.py: -------------------------------------------------------------------------------- 1 | import copy 2 | _base_ = '../../base.py' 3 | # model settings 4 | model = dict( 5 | type='BYOL', 6 | pretrained=None, 7 | base_momentum=0.996, 8 | backbone=dict( 9 | type='ResNet', 10 | depth=50, 11 | in_channels=3, 12 | out_indices=[4], # 0: conv-1, x: stage-x 13 | norm_cfg=dict(type='BN')), 14 | neck=dict( 15 | type='NonLinearNeckV2', 16 | in_channels=2048, 17 | hid_channels=4096, 18 | out_channels=256, 19 | with_avg_pool=True), 20 | head=dict(type='LatentPredictHead', 21 | size_average=True, 22 | predictor=dict(type='NonLinearNeckV2', 23 | in_channels=256, hid_channels=4096, 24 | out_channels=256, with_avg_pool=False))) 25 | # dataset settings 26 | data_source_cfg = dict( 27 | type='ImageNet', 28 | memcached=True, 29 | mclient_path='/mnt/lustre/share/memcached_client') 30 | data_train_list = 'data/imagenet/meta/train.txt' 31 | data_train_root = 'data/imagenet/train' 32 | dataset_type = 'BYOLDataset' 33 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 34 | train_pipeline = [ 35 | dict(type='RandomResizedCrop', size=224, interpolation=3), # bicubic 36 | dict(type='RandomHorizontalFlip'), 37 | dict( 38 | type='RandomAppliedTrans', 39 | transforms=[ 40 | dict( 41 | type='ColorJitter', 42 | brightness=0.4, 43 | contrast=0.4, 44 | saturation=0.2, 45 | hue=0.1) 46 | ], 47 | p=0.8), 48 | dict(type='RandomGrayscale', p=0.2), 49 | dict( 50 | type='RandomAppliedTrans', 51 | transforms=[ 52 | dict( 53 | type='GaussianBlur', 54 | sigma_min=0.1, 55 | sigma_max=2.0) 56 | ], 57 | p=1.), 58 | dict(type='RandomAppliedTrans', 59 | transforms=[dict(type='Solarization')], p=0.), 60 | ] 61 | # prefetch 62 | prefetch = False 63 | if not prefetch: 64 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 65 | train_pipeline1 = copy.deepcopy(train_pipeline) 66 | train_pipeline2 = copy.deepcopy(train_pipeline) 67 | train_pipeline2[4]['p'] = 0.1 # gaussian blur 68 | train_pipeline2[5]['p'] = 0.2 # solarization 69 | 70 | data = dict( 71 | imgs_per_gpu=32, # total 32*8=256 72 | workers_per_gpu=4, 73 | train=dict( 74 | type=dataset_type, 75 | data_source=dict( 76 | list_file=data_train_list, root=data_train_root, 77 | **data_source_cfg), 78 | pipeline1=train_pipeline1, 79 | pipeline2=train_pipeline2, 80 | prefetch=prefetch, 81 | )) 82 | # additional hooks 83 | custom_hooks = [ 84 | dict(type='BYOLHook', end_momentum=1.) 85 | ] 86 | # optimizer 87 | optimizer = dict(type='LARS', lr=0.2, weight_decay=0.0000015, momentum=0.9, 88 | paramwise_options={ 89 | '(bn|gn)(\d+)?.(weight|bias)': dict(weight_decay=0., lars_exclude=True), 90 | 'bias': dict(weight_decay=0., lars_exclude=True)}) 91 | # learning policy 92 | lr_config = dict( 93 | policy='CosineAnnealing', 94 | min_lr=0., 95 | warmup='linear', 96 | warmup_iters=2, 97 | warmup_ratio=0.0001, # cannot be 0 98 | warmup_by_epoch=True) 99 | checkpoint_config = dict(interval=10) 100 | # runtime settings 101 | total_epochs = 200 102 | -------------------------------------------------------------------------------- /configs/selfsup/deepcluster/r50.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | num_classes = 10000 4 | model = dict( 5 | type='DeepCluster', 6 | pretrained=None, 7 | with_sobel=True, 8 | backbone=dict( 9 | type='ResNet', 10 | depth=50, 11 | in_channels=2, 12 | out_indices=[4], # 0: conv-1, x: stage-x 13 | norm_cfg=dict(type='BN')), 14 | neck=dict(type='AvgPoolNeck'), 15 | head=dict( 16 | type='ClsHead', 17 | with_avg_pool=False, # already has avgpool in the neck 18 | in_channels=2048, 19 | num_classes=num_classes)) 20 | # dataset settings 21 | data_source_cfg = dict( 22 | type='ImageNet', 23 | memcached=True, 24 | mclient_path='/mnt/lustre/share/memcached_client') 25 | data_train_list = 'data/imagenet/meta/train.txt' 26 | data_train_root = 'data/imagenet/train' 27 | dataset_type = 'DeepClusterDataset' 28 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 29 | train_pipeline = [ 30 | dict(type='RandomResizedCrop', size=224), 31 | dict(type='RandomHorizontalFlip'), 32 | dict(type='RandomRotation', degrees=2), 33 | dict( 34 | type='ColorJitter', 35 | brightness=0.4, 36 | contrast=0.4, 37 | saturation=1.0, 38 | hue=0.5), 39 | dict(type='RandomGrayscale', p=0.2), 40 | dict(type='ToTensor'), 41 | dict(type='Normalize', **img_norm_cfg), 42 | ] 43 | extract_pipeline = [ 44 | dict(type='Resize', size=256), 45 | dict(type='CenterCrop', size=224), 46 | dict(type='ToTensor'), 47 | dict(type='Normalize', **img_norm_cfg), 48 | ] 49 | data = dict( 50 | imgs_per_gpu=64, 51 | workers_per_gpu=4, 52 | train=dict( 53 | type=dataset_type, 54 | data_source=dict( 55 | list_file=data_train_list, root=data_train_root, 56 | **data_source_cfg), 57 | pipeline=train_pipeline)) 58 | # additional hooks 59 | custom_hooks = [ 60 | dict( 61 | type='DeepClusterHook', 62 | extractor=dict( 63 | imgs_per_gpu=128, 64 | workers_per_gpu=8, 65 | dataset=dict( 66 | type=dataset_type, 67 | data_source=dict( 68 | list_file=data_train_list, 69 | root=data_train_root, 70 | **data_source_cfg), 71 | pipeline=extract_pipeline)), 72 | clustering=dict(type='Kmeans', k=num_classes, pca_dim=256), 73 | unif_sampling=True, 74 | reweight=False, 75 | reweight_pow=0.5, 76 | initial=True, # call initially 77 | interval=1) 78 | ] 79 | # optimizer 80 | optimizer = dict( 81 | type='SGD', lr=0.1, momentum=0.9, weight_decay=0.00001, 82 | nesterov=False, 83 | paramwise_options={'\Ahead.': dict(momentum=0.)}) 84 | # learning policy 85 | lr_config = dict(policy='step', step=[400]) 86 | checkpoint_config = dict(interval=10) 87 | # runtime settings 88 | total_epochs = 200 89 | -------------------------------------------------------------------------------- /configs/selfsup/deepcluster/r50_withoutsobel.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | num_classes = 10000 4 | model = dict( 5 | type='DeepCluster', 6 | pretrained=None, 7 | with_sobel=False, 8 | backbone=dict( 9 | type='ResNet', 10 | depth=50, 11 | in_channels=3, 12 | out_indices=[4], # 0: conv-1, x: stage-x 13 | norm_cfg=dict(type='BN')), 14 | neck=dict(type='AvgPoolNeck'), 15 | head=dict( 16 | type='ClsHead', 17 | with_avg_pool=False, # already has avgpool in the neck 18 | in_channels=2048, 19 | num_classes=num_classes)) 20 | # dataset settings 21 | data_source_cfg = dict( 22 | type='ImageNet', 23 | memcached=True, 24 | mclient_path='/mnt/lustre/share/memcached_client') 25 | data_train_list = 'data/imagenet/meta/train.txt' 26 | data_train_root = 'data/imagenet/jpeg/train' 27 | dataset_type = 'DeepClusterDataset' 28 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 29 | train_pipeline = [ 30 | dict(type='RandomResizedCrop', size=224), 31 | dict(type='RandomHorizontalFlip'), 32 | dict(type='RandomRotation', degrees=2), 33 | dict( 34 | type='ColorJitter', 35 | brightness=0.4, 36 | contrast=0.4, 37 | saturation=1.0, 38 | hue=0.5), 39 | dict(type='RandomGrayscale', p=0.2), 40 | dict(type='ToTensor'), 41 | dict(type='Normalize', **img_norm_cfg), 42 | ] 43 | extract_pipeline = [ 44 | dict(type='Resize', size=256), 45 | dict(type='CenterCrop', size=224), 46 | dict(type='ToTensor'), 47 | dict(type='Normalize', **img_norm_cfg), 48 | ] 49 | data = dict( 50 | imgs_per_gpu=64, # 32 51 | workers_per_gpu=4, 52 | train=dict( 53 | type=dataset_type, 54 | data_source=dict( 55 | list_file=data_train_list, root=data_train_root, 56 | **data_source_cfg), 57 | pipeline=train_pipeline)) 58 | # additional hooks 59 | custom_hooks = [ 60 | dict( 61 | type='DeepClusterHook', 62 | extractor=dict( 63 | imgs_per_gpu=128, 64 | workers_per_gpu=8, 65 | dataset=dict( 66 | type=dataset_type, 67 | data_source=dict( 68 | list_file=data_train_list, 69 | root=data_train_root, 70 | **data_source_cfg), 71 | pipeline=extract_pipeline)), 72 | clustering=dict(type='Kmeans', k=num_classes, pca_dim=256), 73 | unif_sampling=True, 74 | reweight=False, 75 | reweight_pow=0.5, 76 | initial=True, # call initially 77 | interval=1) 78 | ] 79 | # optimizer 80 | optimizer = dict( 81 | type='SGD', lr=0.3, momentum=0.9, weight_decay=0.00001, 82 | nesterov=False, 83 | paramwise_options={'\Ahead.': dict(momentum=0.)}) 84 | # learning policy 85 | lr_config = dict(policy='step', step=[400]) 86 | checkpoint_config = dict(interval=10) 87 | # runtime settings 88 | total_epochs = 200 89 | -------------------------------------------------------------------------------- /configs/selfsup/moco/r50_v1.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='MOCO', 5 | pretrained=None, 6 | queue_len=65536, 7 | feat_dim=128, 8 | momentum=0.999, 9 | backbone=dict( 10 | type='ResNet', 11 | depth=50, 12 | in_channels=3, 13 | out_indices=[4], # 0: conv-1, x: stage-x 14 | norm_cfg=dict(type='BN')), 15 | neck=dict( 16 | type='LinearNeck', 17 | in_channels=2048, 18 | out_channels=128, 19 | with_avg_pool=True), 20 | head=dict(type='ContrastiveHead', temperature=0.07)) 21 | # dataset settings 22 | data_source_cfg = dict( 23 | type='ImageNet', 24 | memcached=True, 25 | mclient_path='/mnt/lustre/share/memcached_client') 26 | data_train_list = 'data/imagenet/meta/train.txt' 27 | data_train_root = 'data/imagenet/train' 28 | dataset_type = 'ContrastiveDataset' 29 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 30 | train_pipeline = [ 31 | dict(type='RandomResizedCrop', size=224, scale=(0.2, 1.)), 32 | dict(type='RandomGrayscale', p=0.2), 33 | dict( 34 | type='ColorJitter', 35 | brightness=0.4, 36 | contrast=0.4, 37 | saturation=0.4, 38 | hue=0.4), 39 | dict(type='RandomHorizontalFlip'), 40 | ] 41 | # prefetch 42 | prefetch = False 43 | if not prefetch: 44 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 45 | 46 | data = dict( 47 | imgs_per_gpu=32, # total 32*8=256 48 | workers_per_gpu=4, 49 | drop_last=True, 50 | train=dict( 51 | type=dataset_type, 52 | data_source=dict( 53 | list_file=data_train_list, root=data_train_root, 54 | **data_source_cfg), 55 | pipeline=train_pipeline, 56 | prefetch=prefetch, 57 | )) 58 | # optimizer 59 | optimizer = dict(type='SGD', lr=0.03, weight_decay=0.0001, momentum=0.9) 60 | # learning policy 61 | lr_config = dict(policy='step', step=[120, 160]) 62 | checkpoint_config = dict(interval=10) 63 | # runtime settings 64 | total_epochs = 200 65 | -------------------------------------------------------------------------------- /configs/selfsup/moco/r50_v1_fp16.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='MOCO', 5 | pretrained=None, 6 | queue_len=65536, 7 | feat_dim=128, 8 | momentum=0.999, 9 | backbone=dict( 10 | type='ResNet', 11 | depth=50, 12 | in_channels=3, 13 | out_indices=[4], # 0: conv-1, x: stage-x 14 | norm_cfg=dict(type='BN')), 15 | neck=dict( 16 | type='LinearNeck', 17 | in_channels=2048, 18 | out_channels=128, 19 | with_avg_pool=True), 20 | head=dict(type='ContrastiveHead', temperature=0.07)) 21 | # dataset settings 22 | data_source_cfg = dict( 23 | type='ImageNet', 24 | memcached=True, 25 | mclient_path='/mnt/lustre/share/memcached_client') 26 | data_train_list = 'data/imagenet/meta/train.txt' 27 | data_train_root = 'data/imagenet/train' 28 | dataset_type = 'ContrastiveDataset' 29 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 30 | train_pipeline = [ 31 | dict(type='RandomResizedCrop', size=224, scale=(0.2, 1.)), 32 | dict(type='RandomGrayscale', p=0.2), 33 | dict( 34 | type='ColorJitter', 35 | brightness=0.4, 36 | contrast=0.4, 37 | saturation=0.4, 38 | hue=0.4), 39 | dict(type='RandomHorizontalFlip'), 40 | ] 41 | # prefetch 42 | prefetch = False 43 | if not prefetch: 44 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 45 | 46 | data = dict( 47 | imgs_per_gpu=32, # total 32*8=256 48 | workers_per_gpu=4, 49 | drop_last=True, 50 | train=dict( 51 | type=dataset_type, 52 | data_source=dict( 53 | list_file=data_train_list, root=data_train_root, 54 | **data_source_cfg), 55 | pipeline=train_pipeline, 56 | prefetch=prefetch, 57 | )) 58 | # optimizer 59 | optimizer = dict(type='SGD', lr=0.03, weight_decay=0.0001, momentum=0.9) 60 | # learning policy 61 | lr_config = dict(policy='step', step=[120, 160]) 62 | checkpoint_config = dict(interval=10) 63 | # runtime settings 64 | total_epochs = 200 65 | # apex 66 | use_fp16 = True 67 | optimizer_config = dict(use_fp16=use_fp16) # grad_clip, coalesce, bucket_size_mb, fp16 68 | 69 | -------------------------------------------------------------------------------- /configs/selfsup/moco/r50_v2.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='MOCO', 5 | pretrained=None, 6 | queue_len=65536, 7 | feat_dim=128, 8 | momentum=0.999, 9 | backbone=dict( 10 | type='ResNet', 11 | depth=50, 12 | in_channels=3, 13 | out_indices=[4], # 0: conv-1, x: stage-x 14 | norm_cfg=dict(type='BN')), 15 | neck=dict( 16 | type='NonLinearNeckV1', 17 | in_channels=2048, 18 | hid_channels=2048, 19 | out_channels=128, 20 | with_avg_pool=True), 21 | head=dict(type='ContrastiveHead', temperature=0.2)) 22 | # dataset settings 23 | data_source_cfg = dict( 24 | type='ImageNet', 25 | memcached=True, 26 | mclient_path='/mnt/lustre/share/memcached_client') 27 | data_train_list = 'data/imagenet/meta/train.txt' 28 | data_train_root = 'data/imagenet/train' 29 | dataset_type = 'ContrastiveDataset' 30 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 31 | train_pipeline = [ 32 | dict(type='RandomResizedCrop', size=224, scale=(0.2, 1.)), 33 | dict( 34 | type='RandomAppliedTrans', 35 | transforms=[ 36 | dict( 37 | type='ColorJitter', 38 | brightness=0.4, 39 | contrast=0.4, 40 | saturation=0.4, 41 | hue=0.1) 42 | ], 43 | p=0.8), 44 | dict(type='RandomGrayscale', p=0.2), 45 | dict( 46 | type='RandomAppliedTrans', 47 | transforms=[ 48 | dict( 49 | type='GaussianBlur', 50 | sigma_min=0.1, 51 | sigma_max=2.0) 52 | ], 53 | p=0.5), 54 | dict(type='RandomHorizontalFlip'), 55 | ] 56 | # prefetch 57 | prefetch = False 58 | if not prefetch: 59 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 60 | 61 | data = dict( 62 | imgs_per_gpu=32, # total 32*8=256 63 | workers_per_gpu=4, 64 | drop_last=True, 65 | train=dict( 66 | type=dataset_type, 67 | data_source=dict( 68 | list_file=data_train_list, root=data_train_root, 69 | **data_source_cfg), 70 | pipeline=train_pipeline, 71 | prefetch=prefetch, 72 | )) 73 | # optimizer 74 | optimizer = dict(type='SGD', lr=0.03, weight_decay=0.0001, momentum=0.9) 75 | # learning policy 76 | lr_config = dict(policy='CosineAnnealing', min_lr=0.) 77 | checkpoint_config = dict(interval=10) 78 | # runtime settings 79 | total_epochs = 200 80 | -------------------------------------------------------------------------------- /configs/selfsup/moco/r50_v2_simclr_neck.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='MOCO', 5 | pretrained=None, 6 | queue_len=65536, 7 | feat_dim=128, 8 | momentum=0.999, 9 | backbone=dict( 10 | type='ResNet', 11 | depth=50, 12 | in_channels=3, 13 | out_indices=[4], # 0: conv-1, x: stage-x 14 | norm_cfg=dict(type='BN')), 15 | neck=dict( 16 | type='NonLinearNeckSimCLR', # SimCLR non-linear neck 17 | in_channels=2048, 18 | hid_channels=2048, 19 | out_channels=128, 20 | num_layers=2, 21 | with_avg_pool=True), 22 | head=dict(type='ContrastiveHead', temperature=0.2)) 23 | # dataset settings 24 | data_source_cfg = dict( 25 | type='ImageNet', 26 | memcached=True, 27 | mclient_path='/mnt/lustre/share/memcached_client') 28 | data_train_list = 'data/imagenet/meta/train.txt' 29 | data_train_root = 'data/imagenet/train' 30 | dataset_type = 'ContrastiveDataset' 31 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 32 | train_pipeline = [ 33 | dict(type='RandomResizedCrop', size=224, scale=(0.2, 1.)), 34 | dict( 35 | type='RandomAppliedTrans', 36 | transforms=[ 37 | dict( 38 | type='ColorJitter', 39 | brightness=0.4, 40 | contrast=0.4, 41 | saturation=0.4, 42 | hue=0.1) 43 | ], 44 | p=0.8), 45 | dict(type='RandomGrayscale', p=0.2), 46 | dict( 47 | type='RandomAppliedTrans', 48 | transforms=[ 49 | dict( 50 | type='GaussianBlur', 51 | sigma_min=0.1, 52 | sigma_max=2.0) 53 | ], 54 | p=0.5), 55 | dict(type='RandomHorizontalFlip'), 56 | ] 57 | 58 | # prefetch 59 | prefetch = False 60 | if not prefetch: 61 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 62 | 63 | data = dict( 64 | imgs_per_gpu=32, # total 32*8=256 65 | workers_per_gpu=4, 66 | drop_last=True, 67 | train=dict( 68 | type=dataset_type, 69 | data_source=dict( 70 | list_file=data_train_list, root=data_train_root, 71 | **data_source_cfg), 72 | pipeline=train_pipeline, 73 | prefetch=prefetch, 74 | )) 75 | # optimizer 76 | optimizer = dict(type='SGD', lr=0.03, weight_decay=0.0001, momentum=0.9) 77 | # learning policy 78 | lr_config = dict(policy='CosineAnnealing', min_lr=0.) 79 | checkpoint_config = dict(interval=10) 80 | # runtime settings 81 | total_epochs = 200 82 | -------------------------------------------------------------------------------- /configs/selfsup/npid/r50.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='NPID', 5 | neg_num=65536, 6 | pretrained=None, 7 | backbone=dict( 8 | type='ResNet', 9 | depth=50, 10 | in_channels=3, 11 | out_indices=[4], # 0: conv-1, x: stage-x 12 | norm_cfg=dict(type='SyncBN')), 13 | neck=dict( 14 | type='LinearNeck', 15 | in_channels=2048, 16 | out_channels=128, 17 | with_avg_pool=True), 18 | head=dict(type='ContrastiveHead', temperature=0.07), 19 | memory_bank=dict( 20 | type='SimpleMemory', length=1281167, feat_dim=128, momentum=0.5)) 21 | # dataset settings 22 | data_source_cfg = dict( 23 | type='ImageNet', 24 | memcached=True, 25 | mclient_path='/mnt/lustre/share/memcached_client') 26 | data_train_list = 'data/imagenet/meta/train.txt' 27 | data_train_root = 'data/imagenet/train' 28 | dataset_type = 'NPIDDataset' 29 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 30 | train_pipeline = [ 31 | dict(type='RandomResizedCrop', size=224, scale=(0.2, 1.)), 32 | dict(type='RandomGrayscale', p=0.2), 33 | dict( 34 | type='ColorJitter', 35 | brightness=0.4, 36 | contrast=0.4, 37 | saturation=0.4, 38 | hue=0.4), 39 | dict(type='RandomHorizontalFlip'), 40 | dict(type='ToTensor'), 41 | dict(type='Normalize', **img_norm_cfg), 42 | ] 43 | test_pipeline = [ 44 | dict(type='Resize', size=256), 45 | dict(type='CenterCrop', size=224), 46 | dict(type='ToTensor'), 47 | dict(type='Normalize', **img_norm_cfg), 48 | ] 49 | data = dict( 50 | imgs_per_gpu=32, # total 32*8 51 | workers_per_gpu=4, 52 | train=dict( 53 | type=dataset_type, 54 | data_source=dict( 55 | list_file=data_train_list, root=data_train_root, 56 | **data_source_cfg), 57 | pipeline=train_pipeline)) 58 | # optimizer 59 | optimizer = dict( 60 | type='SGD', lr=0.03, weight_decay=0.0001, momentum=0.9, nesterov=False) 61 | # learning policy 62 | lr_config = dict(policy='step', step=[120, 160]) 63 | checkpoint_config = dict(interval=10) 64 | # runtime settings 65 | total_epochs = 200 66 | -------------------------------------------------------------------------------- /configs/selfsup/npid/r50_ensure_neg.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='NPID', 5 | neg_num=65536, 6 | ensure_neg=True, 7 | pretrained=None, 8 | backbone=dict( 9 | type='ResNet', 10 | depth=50, 11 | in_channels=3, 12 | out_indices=[4], # 0: conv-1, x: stage-x 13 | norm_cfg=dict(type='SyncBN')), 14 | neck=dict( 15 | type='LinearNeck', 16 | in_channels=2048, 17 | out_channels=128, 18 | with_avg_pool=True), 19 | head=dict(type='ContrastiveHead', temperature=0.07), 20 | memory_bank=dict( 21 | type='SimpleMemory', length=1281167, feat_dim=128, momentum=0.5)) 22 | # dataset settings 23 | data_source_cfg = dict( 24 | type='ImageNet', 25 | memcached=True, 26 | mclient_path='/mnt/lustre/share/memcached_client') 27 | data_train_list = 'data/imagenet/meta/train.txt' 28 | data_train_root = 'data/imagenet/train' 29 | dataset_type = 'NPIDDataset' 30 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 31 | train_pipeline = [ 32 | dict(type='RandomResizedCrop', size=224, scale=(0.2, 1.)), 33 | dict(type='RandomGrayscale', p=0.2), 34 | dict( 35 | type='ColorJitter', 36 | brightness=0.4, 37 | contrast=0.4, 38 | saturation=0.4, 39 | hue=0.4), 40 | dict(type='RandomHorizontalFlip'), 41 | dict(type='ToTensor'), 42 | dict(type='Normalize', **img_norm_cfg), 43 | ] 44 | test_pipeline = [ 45 | dict(type='Resize', size=256), 46 | dict(type='CenterCrop', size=224), 47 | dict(type='ToTensor'), 48 | dict(type='Normalize', **img_norm_cfg), 49 | ] 50 | data = dict( 51 | imgs_per_gpu=32, # total 32*8 52 | workers_per_gpu=4, 53 | train=dict( 54 | type=dataset_type, 55 | data_source=dict( 56 | list_file=data_train_list, root=data_train_root, 57 | **data_source_cfg), 58 | pipeline=train_pipeline)) 59 | # optimizer 60 | optimizer = dict( 61 | type='SGD', lr=0.03, weight_decay=0.0001, momentum=0.9, nesterov=False) 62 | # learning policy 63 | lr_config = dict(policy='step', step=[120, 160]) 64 | checkpoint_config = dict(interval=10) 65 | # runtime settings 66 | total_epochs = 200 67 | -------------------------------------------------------------------------------- /configs/selfsup/odc/r50_v1.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | num_classes = 10000 4 | model = dict( 5 | type='ODC', 6 | pretrained=None, 7 | with_sobel=False, 8 | backbone=dict( 9 | type='ResNet', 10 | depth=50, 11 | in_channels=3, 12 | out_indices=[4], # 0: conv-1, x: stage-x 13 | norm_cfg=dict(type='SyncBN')), 14 | neck=dict( 15 | type='NonLinearNeckV0', 16 | in_channels=2048, 17 | hid_channels=512, 18 | out_channels=256, 19 | with_avg_pool=True), 20 | head=dict( 21 | type='ClsHead', 22 | with_avg_pool=False, 23 | in_channels=256, 24 | num_classes=num_classes), 25 | memory_bank=dict( 26 | type='ODCMemory', 27 | length=1281167, 28 | feat_dim=256, 29 | momentum=0.5, 30 | num_classes=num_classes, 31 | min_cluster=20, 32 | debug=False)) 33 | # dataset settings 34 | data_source_cfg = dict( 35 | type='ImageNet', 36 | memcached=True, 37 | mclient_path='/mnt/lustre/share/memcached_client') 38 | data_train_list = 'data/imagenet/meta/train.txt' 39 | data_train_root = 'data/imagenet/train' 40 | dataset_type = 'DeepClusterDataset' 41 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 42 | train_pipeline = [ 43 | dict(type='RandomResizedCrop', size=224), 44 | dict(type='RandomHorizontalFlip'), 45 | dict(type='RandomRotation', degrees=2), 46 | dict( 47 | type='ColorJitter', 48 | brightness=0.4, 49 | contrast=0.4, 50 | saturation=1.0, 51 | hue=0.5), 52 | dict(type='RandomGrayscale', p=0.2), 53 | dict(type='ToTensor'), 54 | dict(type='Normalize', **img_norm_cfg), 55 | ] 56 | extract_pipeline = [ 57 | dict(type='Resize', size=256), 58 | dict(type='CenterCrop', size=224), 59 | dict(type='ToTensor'), 60 | dict(type='Normalize', **img_norm_cfg), 61 | ] 62 | data = dict( 63 | imgs_per_gpu=64, # 64*8 64 | sampling_replace=True, 65 | workers_per_gpu=2, 66 | train=dict( 67 | type=dataset_type, 68 | data_source=dict( 69 | list_file=data_train_list, root=data_train_root, 70 | **data_source_cfg), 71 | pipeline=train_pipeline)) 72 | # additional hooks 73 | custom_hooks = [ 74 | dict( 75 | type='DeepClusterHook', 76 | extractor=dict( 77 | imgs_per_gpu=128, 78 | workers_per_gpu=8, 79 | dataset=dict( 80 | type=dataset_type, 81 | data_source=dict( 82 | list_file=data_train_list, 83 | root=data_train_root, 84 | **data_source_cfg), 85 | pipeline=extract_pipeline)), 86 | clustering=dict(type='Kmeans', k=num_classes, pca_dim=-1), # no pca 87 | unif_sampling=False, 88 | reweight=True, 89 | reweight_pow=0.5, 90 | init_memory=True, 91 | initial=True, # call initially 92 | interval=9999999999), # initial only 93 | dict( 94 | type='ODCHook', 95 | centroids_update_interval=10, # iter 96 | deal_with_small_clusters_interval=1, 97 | evaluate_interval=50, 98 | reweight=True, 99 | reweight_pow=0.5) 100 | ] 101 | # optimizer 102 | optimizer = dict( 103 | type='SGD', lr=0.06, momentum=0.9, weight_decay=0.00001, 104 | nesterov=False, 105 | paramwise_options={'\Ahead.': dict(momentum=0.)}) 106 | # learning policy 107 | lr_config = dict(policy='step', step=[400], gamma=0.4) 108 | checkpoint_config = dict(interval=10) 109 | # runtime settings 110 | total_epochs = 440 111 | -------------------------------------------------------------------------------- /configs/selfsup/regionCL-M/regionCLM-r50.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='RegionCLM', 5 | pretrained=None, 6 | queue_len=65536, 7 | feat_dim=128, 8 | momentum=0.999, 9 | loss_cutmix=1.0, 10 | cutmix_alpha=1.0, 11 | cutMixUpper=6, 12 | cutMixLower=3, 13 | backbone=dict( 14 | type='ResNet', 15 | depth=50, 16 | in_channels=3, 17 | out_indices=[4], # 0: conv-1, x: stage-x 18 | norm_cfg=dict(type='BN')), 19 | neck=dict( 20 | type='RegionCLNonLinearNeckV1', 21 | in_channels=2048, 22 | hid_channels=2048, 23 | out_channels=128, 24 | with_avg_pool=True), 25 | head=dict(type='ContrastiveHead', temperature=0.2)) 26 | # dataset settings 27 | data_source_cfg = dict( 28 | type='ImageNet', 29 | memcached=True, 30 | mclient_path='/mnt/lustre/share/memcached_client') 31 | data_train_list = 'data/imagenet/meta/train.txt' 32 | data_train_root = 'data/imagenet/train' 33 | dataset_type = 'ContrastiveDataset' 34 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 35 | train_pipeline = [ 36 | dict(type='RandomResizedCrop', size=224, scale=(0.2, 1.)), 37 | dict( 38 | type='RandomAppliedTrans', 39 | transforms=[ 40 | dict( 41 | type='ColorJitter', 42 | brightness=0.4, 43 | contrast=0.4, 44 | saturation=0.4, 45 | hue=0.1) 46 | ], 47 | p=0.8), 48 | dict(type='RandomGrayscale', p=0.2), 49 | dict( 50 | type='RandomAppliedTrans', 51 | transforms=[ 52 | dict( 53 | type='GaussianBlur', 54 | sigma_min=0.1, 55 | sigma_max=2.0) 56 | ], 57 | p=0.5), 58 | dict(type='RandomHorizontalFlip'), 59 | ] 60 | # prefetch 61 | prefetch = False 62 | if not prefetch: 63 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 64 | data = dict( 65 | imgs_per_gpu=128, # total 128*8=1024 66 | workers_per_gpu=10, 67 | drop_last=True, 68 | train=dict( 69 | type=dataset_type, 70 | data_source=dict( 71 | list_file=data_train_list, root=data_train_root, 72 | **data_source_cfg), 73 | pipeline=train_pipeline, 74 | prefetch=prefetch, 75 | )) 76 | # optimizer 77 | optimizer = dict(type='SGD', lr=0.15, weight_decay=0.0001, momentum=0.9) 78 | lr_config = dict(policy='CosineAnnealing', min_lr=0.) 79 | checkpoint_config = dict(interval=10) 80 | # runtime settings 81 | total_epochs = 200 82 | find_unused_parameters=True 83 | -------------------------------------------------------------------------------- /configs/selfsup/relative_loc/r50.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='RelativeLoc', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | in_channels=3, 10 | out_indices=[4], # 0: conv-1, x: stage-x 11 | norm_cfg=dict(type='BN')), 12 | neck=dict( 13 | type='RelativeLocNeck', 14 | in_channels=2048, 15 | out_channels=4096, 16 | with_avg_pool=True), 17 | head=dict( 18 | type='ClsHead', 19 | with_avg_pool=False, 20 | in_channels=4096, 21 | num_classes=8)) 22 | # dataset settings 23 | data_source_cfg = dict( 24 | type='ImageNet', 25 | memcached=True, 26 | mclient_path='/mnt/lustre/share/memcached_client') 27 | data_train_list = 'data/imagenet/meta/train.txt' 28 | data_train_root = 'data/imagenet/train' 29 | data_test_list = 'data/imagenet/meta/val.txt' 30 | data_test_root = 'data/imagenet/val' 31 | dataset_type = 'RelativeLocDataset' 32 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 33 | train_pipeline = [ 34 | dict(type='Resize', size=292), 35 | dict(type='RandomCrop', size=255), 36 | dict(type='RandomGrayscale', p=0.66), 37 | ] 38 | test_pipeline = [ 39 | dict(type='Resize', size=292), 40 | dict(type='CenterCrop', size=255), 41 | ] 42 | format_pipeline = [ 43 | dict(type='ToTensor'), 44 | dict(type='Normalize', **img_norm_cfg), 45 | ] 46 | data = dict( 47 | imgs_per_gpu=64, # 64 x 8 = 512 48 | workers_per_gpu=2, 49 | train=dict( 50 | type=dataset_type, 51 | data_source=dict( 52 | list_file=data_train_list, root=data_train_root, 53 | **data_source_cfg), 54 | pipeline=train_pipeline, 55 | format_pipeline=format_pipeline), 56 | val=dict( 57 | type=dataset_type, 58 | data_source=dict( 59 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 60 | pipeline=test_pipeline, 61 | format_pipeline=format_pipeline)) 62 | # optimizer 63 | optimizer = dict( 64 | type='SGD', lr=0.2, momentum=0.9, weight_decay=0.0001, 65 | nesterov=False, 66 | paramwise_options={ 67 | '\Aneck.': dict(weight_decay=0.0005), 68 | '\Ahead.': dict(weight_decay=0.0005)}) 69 | # learning policy 70 | lr_config = dict( 71 | policy='step', 72 | step=[30, 50], 73 | warmup='linear', 74 | warmup_iters=5, # 5 ep 75 | warmup_ratio=0.1, 76 | warmup_by_epoch=True) 77 | checkpoint_config = dict(interval=10) 78 | # runtime settings 79 | total_epochs = 70 80 | -------------------------------------------------------------------------------- /configs/selfsup/rotation_pred/r50.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='RotationPred', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | in_channels=3, 10 | out_indices=[4], # 0: conv-1, x: stage-x 11 | norm_cfg=dict(type='SyncBN')), 12 | head=dict( 13 | type='ClsHead', with_avg_pool=True, in_channels=2048, num_classes=4)) 14 | # dataset settings 15 | data_source_cfg = dict( 16 | type='ImageNet', 17 | memcached=True, 18 | mclient_path='/mnt/lustre/share/memcached_client') 19 | data_train_list = 'data/imagenet/meta/train.txt' 20 | data_train_root = 'data/imagenet/train' 21 | data_test_list = 'data/imagenet/meta/val.txt' 22 | data_test_root = 'data/imagenet/val' 23 | dataset_type = 'RotationPredDataset' 24 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 25 | train_pipeline = [ 26 | dict(type='RandomResizedCrop', size=224), 27 | dict(type='RandomHorizontalFlip'), 28 | dict(type='ToTensor'), 29 | dict(type='Normalize', **img_norm_cfg), 30 | ] 31 | test_pipeline = [ 32 | dict(type='Resize', size=256), 33 | dict(type='CenterCrop', size=224), 34 | dict(type='ToTensor'), 35 | dict(type='Normalize', **img_norm_cfg), 36 | ] 37 | data = dict( 38 | imgs_per_gpu=16, # (16*4) x 8 = 512 39 | workers_per_gpu=2, 40 | train=dict( 41 | type=dataset_type, 42 | data_source=dict( 43 | list_file=data_train_list, root=data_train_root, 44 | **data_source_cfg), 45 | pipeline=train_pipeline), 46 | val=dict( 47 | type=dataset_type, 48 | data_source=dict( 49 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 50 | pipeline=test_pipeline)) 51 | # optimizer 52 | optimizer = dict( 53 | type='SGD', lr=0.2, momentum=0.9, weight_decay=0.0001, nesterov=False) 54 | # learning policy 55 | lr_config = dict( 56 | policy='step', 57 | step=[30, 50], 58 | warmup='linear', 59 | warmup_iters=5, # 5 ep 60 | warmup_ratio=0.1, 61 | warmup_by_epoch=True) 62 | checkpoint_config = dict(interval=10) 63 | # runtime settings 64 | total_epochs = 70 65 | -------------------------------------------------------------------------------- /configs/selfsup/rotation_pred/r50_fp16.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='RotationPred', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | in_channels=3, 10 | out_indices=[4], # 0: conv-1, x: stage-x 11 | norm_cfg=dict(type='SyncBN')), 12 | head=dict( 13 | type='ClsHead', with_avg_pool=True, in_channels=2048, num_classes=4)) 14 | # dataset settings 15 | data_source_cfg = dict( 16 | type='ImageNet', 17 | memcached=True, 18 | mclient_path='/mnt/lustre/share/memcached_client') 19 | data_train_list = 'data/imagenet/meta/train.txt' 20 | data_train_root = 'data/imagenet/train' 21 | data_test_list = 'data/imagenet/meta/val.txt' 22 | data_test_root = 'data/imagenet/val' 23 | dataset_type = 'RotationPredDataset' 24 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 25 | train_pipeline = [ 26 | dict(type='RandomResizedCrop', size=224), 27 | dict(type='RandomHorizontalFlip'), 28 | dict(type='ToTensor'), 29 | dict(type='Normalize', **img_norm_cfg), 30 | ] 31 | test_pipeline = [ 32 | dict(type='Resize', size=256), 33 | dict(type='CenterCrop', size=224), 34 | dict(type='ToTensor'), 35 | dict(type='Normalize', **img_norm_cfg), 36 | ] 37 | data = dict( 38 | imgs_per_gpu=16, # (16*4) x 8 = 512 39 | workers_per_gpu=2, 40 | train=dict( 41 | type=dataset_type, 42 | data_source=dict( 43 | list_file=data_train_list, root=data_train_root, 44 | **data_source_cfg), 45 | pipeline=train_pipeline), 46 | val=dict( 47 | type=dataset_type, 48 | data_source=dict( 49 | list_file=data_test_list, root=data_test_root, **data_source_cfg), 50 | pipeline=test_pipeline)) 51 | # optimizer 52 | optimizer = dict( 53 | type='SGD', lr=0.2, momentum=0.9, weight_decay=0.0001, nesterov=False) 54 | # learning policy 55 | lr_config = dict( 56 | policy='step', 57 | step=[30, 50], 58 | warmup='linear', 59 | warmup_iters=5, # 5 ep 60 | warmup_ratio=0.1, 61 | warmup_by_epoch=True) 62 | checkpoint_config = dict(interval=10) 63 | # runtime settings 64 | total_epochs = 70 65 | # apex 66 | use_fp16 = True 67 | optimizer_config = dict(use_fp16=use_fp16) # grad_clip, coalesce, bucket_size_mb, fp16 68 | 69 | -------------------------------------------------------------------------------- /configs/selfsup/simclr/r50_bs256_ep200.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='SimCLR', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | in_channels=3, 10 | out_indices=[4], # 0: conv-1, x: stage-x 11 | norm_cfg=dict(type='SyncBN')), 12 | neck=dict( 13 | type='NonLinearNeckSimCLR', # SimCLR non-linear neck 14 | in_channels=2048, 15 | hid_channels=2048, 16 | out_channels=128, 17 | num_layers=2, 18 | with_avg_pool=True), 19 | head=dict(type='ContrastiveHead', temperature=0.1)) 20 | # dataset settings 21 | data_source_cfg = dict( 22 | type='ImageNet', 23 | memcached=True, 24 | mclient_path='/mnt/lustre/share/memcached_client') 25 | data_train_list = 'data/imagenet/meta/train.txt' 26 | data_train_root = 'data/imagenet/train' 27 | dataset_type = 'ContrastiveDataset' 28 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 29 | train_pipeline = [ 30 | dict(type='RandomResizedCrop', size=224), 31 | dict(type='RandomHorizontalFlip'), 32 | dict( 33 | type='RandomAppliedTrans', 34 | transforms=[ 35 | dict( 36 | type='ColorJitter', 37 | brightness=0.8, 38 | contrast=0.8, 39 | saturation=0.8, 40 | hue=0.2) 41 | ], 42 | p=0.8), 43 | dict(type='RandomGrayscale', p=0.2), 44 | dict( 45 | type='RandomAppliedTrans', 46 | transforms=[ 47 | dict( 48 | type='GaussianBlur', 49 | sigma_min=0.1, 50 | sigma_max=2.0) 51 | ], 52 | p=0.5), 53 | ] 54 | 55 | # prefetch 56 | prefetch = False 57 | if not prefetch: 58 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 59 | 60 | data = dict( 61 | imgs_per_gpu=32, # total 32*8 62 | workers_per_gpu=4, 63 | train=dict( 64 | type=dataset_type, 65 | data_source=dict( 66 | list_file=data_train_list, root=data_train_root, 67 | **data_source_cfg), 68 | pipeline=train_pipeline, 69 | prefetch=prefetch, 70 | )) 71 | # optimizer 72 | optimizer = dict(type='LARS', lr=0.3, weight_decay=0.000001, momentum=0.9, 73 | paramwise_options={ 74 | '(bn|gn)(\d+)?.(weight|bias)': dict(weight_decay=0., lars_exclude=True), 75 | 'bias': dict(weight_decay=0., lars_exclude=True)}) 76 | # learning policy 77 | lr_config = dict( 78 | policy='CosineAnnealing', 79 | min_lr=0., 80 | warmup='linear', 81 | warmup_iters=10, 82 | warmup_ratio=0.0001, 83 | warmup_by_epoch=True) 84 | checkpoint_config = dict(interval=10) 85 | # runtime settings 86 | total_epochs = 200 87 | -------------------------------------------------------------------------------- /configs/selfsup/simclr/r50_bs256_ep200_mocov2_neck.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='SimCLR', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | in_channels=3, 10 | out_indices=[4], # 0: conv-1, x: stage-x 11 | norm_cfg=dict(type='SyncBN')), 12 | neck=dict( 13 | type='NonLinearNeckV1', # simple fc-relu-fc neck in MoCo v2 14 | in_channels=2048, 15 | hid_channels=2048, 16 | out_channels=128, 17 | with_avg_pool=True), 18 | head=dict(type='ContrastiveHead', temperature=0.1)) 19 | # dataset settings 20 | data_source_cfg = dict( 21 | type='ImageNet', 22 | memcached=True, 23 | mclient_path='/mnt/lustre/share/memcached_client') 24 | data_train_list = 'data/imagenet/meta/train.txt' 25 | data_train_root = 'data/imagenet/train' 26 | dataset_type = 'ContrastiveDataset' 27 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 28 | train_pipeline = [ 29 | dict(type='RandomResizedCrop', size=224), 30 | dict(type='RandomHorizontalFlip'), 31 | dict( 32 | type='RandomAppliedTrans', 33 | transforms=[ 34 | dict( 35 | type='ColorJitter', 36 | brightness=0.8, 37 | contrast=0.8, 38 | saturation=0.8, 39 | hue=0.2) 40 | ], 41 | p=0.8), 42 | dict(type='RandomGrayscale', p=0.2), 43 | dict( 44 | type='RandomAppliedTrans', 45 | transforms=[ 46 | dict( 47 | type='GaussianBlur', 48 | sigma_min=0.1, 49 | sigma_max=2.0) 50 | ], 51 | p=0.5), 52 | ] 53 | 54 | # prefetch 55 | prefetch = False 56 | if not prefetch: 57 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 58 | 59 | data = dict( 60 | imgs_per_gpu=32, # total 32*8 61 | workers_per_gpu=4, 62 | train=dict( 63 | type=dataset_type, 64 | data_source=dict( 65 | list_file=data_train_list, root=data_train_root, 66 | **data_source_cfg), 67 | pipeline=train_pipeline, 68 | prefetch=prefetch, 69 | )) 70 | # optimizer 71 | optimizer = dict(type='LARS', lr=0.3, weight_decay=0.000001, momentum=0.9, 72 | paramwise_options={ 73 | '(bn|gn)(\d+)?.(weight|bias)': dict(weight_decay=0., lars_exclude=True), 74 | 'bias': dict(weight_decay=0., lars_exclude=True)}) 75 | # learning policy 76 | lr_config = dict( 77 | policy='CosineAnnealing', 78 | min_lr=0., 79 | warmup='linear', 80 | warmup_iters=10, 81 | warmup_ratio=0.0001, 82 | warmup_by_epoch=True) 83 | checkpoint_config = dict(interval=10) 84 | # runtime settings 85 | total_epochs = 200 86 | -------------------------------------------------------------------------------- /configs/selfsup/simclr/r50_bs512_ep200.py: -------------------------------------------------------------------------------- 1 | _base_ = '../../base.py' 2 | # model settings 3 | model = dict( 4 | type='SimCLR', 5 | pretrained=None, 6 | backbone=dict( 7 | type='ResNet', 8 | depth=50, 9 | in_channels=3, 10 | out_indices=[4], # 0: conv-1, x: stage-x 11 | norm_cfg=dict(type='SyncBN')), 12 | neck=dict( 13 | type='NonLinearNeckSimCLR', # SimCLR non-linear neck 14 | in_channels=2048, 15 | hid_channels=2048, 16 | out_channels=128, 17 | num_layers=2, 18 | with_avg_pool=True), 19 | head=dict(type='ContrastiveHead', temperature=0.1)) 20 | # dataset settings 21 | data_source_cfg = dict( 22 | type='ImageNet', 23 | memcached=True, 24 | mclient_path='/mnt/lustre/share/memcached_client') 25 | data_train_list = 'data/imagenet/meta/train.txt' 26 | data_train_root = 'data/imagenet/train' 27 | dataset_type = 'ContrastiveDataset' 28 | img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 29 | train_pipeline = [ 30 | dict(type='RandomResizedCrop', size=224), 31 | dict(type='RandomHorizontalFlip'), 32 | dict( 33 | type='RandomAppliedTrans', 34 | transforms=[ 35 | dict( 36 | type='ColorJitter', 37 | brightness=0.8, 38 | contrast=0.8, 39 | saturation=0.8, 40 | hue=0.2) 41 | ], 42 | p=0.8), 43 | dict(type='RandomGrayscale', p=0.2), 44 | dict( 45 | type='RandomAppliedTrans', 46 | transforms=[ 47 | dict( 48 | type='GaussianBlur', 49 | sigma_min=0.1, 50 | sigma_max=2.0) 51 | ], 52 | p=0.5), 53 | ] 54 | 55 | # prefetch 56 | prefetch = False 57 | if not prefetch: 58 | train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)]) 59 | 60 | data = dict( 61 | imgs_per_gpu=64, # total 64*8 62 | workers_per_gpu=4, 63 | train=dict( 64 | type=dataset_type, 65 | data_source=dict( 66 | list_file=data_train_list, root=data_train_root, 67 | **data_source_cfg), 68 | pipeline=train_pipeline, 69 | prefetch=prefetch, 70 | )) 71 | # optimizer 72 | optimizer = dict(type='LARS', lr=0.6, weight_decay=0.000001, momentum=0.9, 73 | paramwise_options={ 74 | '(bn|gn)(\d+)?.(weight|bias)': dict(weight_decay=0., lars_exclude=True), 75 | 'bias': dict(weight_decay=0., lars_exclude=True)}) 76 | # learning policy 77 | lr_config = dict( 78 | policy='CosineAnnealing', 79 | min_lr=0., 80 | warmup='linear', 81 | warmup_iters=10, 82 | warmup_ratio=0.0001, 83 | warmup_by_epoch=True) 84 | checkpoint_config = dict(interval=10) 85 | # runtime settings 86 | total_epochs = 200 87 | -------------------------------------------------------------------------------- /demo/NetworkStructure.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Annbless/RegionCL/61bbcabb0a7c8eb91d14bbbd139779e650d150ba/demo/NetworkStructure.png -------------------------------------------------------------------------------- /docs/CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ## Changelog 2 | 3 | ### v0.3.0 (14/10/2020) 4 | 5 | #### Highlight 6 | * Support Mixed Precision Training 7 | * Improvement of GaussianBlur doubles the training speed 8 | * More benchmarking results 9 | 10 | #### Bug Fixes 11 | * Fix bugs in moco v2, now the results are reproducible. 12 | * Fix bugs in byol. 13 | 14 | #### New Features 15 | * Mixed Precision Training 16 | * Improvement of GaussianBlur doubles the training speed of MoCo V2, SimCLR, BYOL 17 | * More benchmarking results, including Places, VOC, COCO 18 | 19 | ### v0.2.0 (26/6/2020) 20 | 21 | #### Highlights 22 | * Support BYOL 23 | * Support semi-supervised benchmarks 24 | 25 | #### Bug Fixes 26 | * Fix hash id in publish_model.py 27 | 28 | #### New Features 29 | 30 | * Support BYOL. 31 | * Separate train and test scripts in linear/semi evaluation. 32 | * Support semi-supevised benchmarks: benchmarks/dist_train_semi.sh. 33 | * Move benchmarks related configs into configs/benchmarks/. 34 | * Provide benchmarking results and model download links. 35 | * Support updating network every several interations. 36 | * Support LARS optimizer with nesterov. 37 | * Support excluding specific parameters from LARS adaptation and weight decay required in SimCLR and BYOL. 38 | -------------------------------------------------------------------------------- /docs/relation.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Annbless/RegionCL/61bbcabb0a7c8eb91d14bbbd139779e650d150ba/docs/relation.jpg -------------------------------------------------------------------------------- /openselfsup/__init__.py: -------------------------------------------------------------------------------- 1 | from .version import __version__, short_version 2 | 3 | __all__ = ['__version__', 'short_version'] 4 | -------------------------------------------------------------------------------- /openselfsup/apis/__init__.py: -------------------------------------------------------------------------------- 1 | from .train import get_root_logger, set_random_seed, train_model 2 | -------------------------------------------------------------------------------- /openselfsup/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_dataset 2 | from .byol import BYOLDataset 3 | from .data_sources import * 4 | from .pipelines import * 5 | from .classification import ClassificationDataset 6 | from .deepcluster import DeepClusterDataset 7 | from .extraction import ExtractDataset 8 | from .npid import NPIDDataset 9 | from .rotation_pred import RotationPredDataset 10 | from .relative_loc import RelativeLocDataset 11 | from .contrastive import ContrastiveDataset 12 | from .dataset_wrappers import ConcatDataset, RepeatDataset 13 | from .loader import DistributedGroupSampler, GroupSampler, build_dataloader 14 | from .registry import DATASETS 15 | -------------------------------------------------------------------------------- /openselfsup/datasets/base.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | import torch 4 | from torch.utils.data import Dataset 5 | 6 | from openselfsup.utils import print_log, build_from_cfg 7 | 8 | from torchvision.transforms import Compose 9 | 10 | from .registry import DATASETS, PIPELINES 11 | from .builder import build_datasource 12 | 13 | 14 | class BaseDataset(Dataset, metaclass=ABCMeta): 15 | """Base dataset. 16 | 17 | Args: 18 | data_source (dict): Data source defined in 19 | `openselfsup.datasets.data_sources`. 20 | pipeline (list[dict]): A list of dict, where each element represents 21 | an operation defined in `oenselfsup.datasets.pipelines`. 22 | """ 23 | 24 | def __init__(self, data_source, pipeline, prefetch=False): 25 | self.data_source = build_datasource(data_source) 26 | pipeline = [build_from_cfg(p, PIPELINES) for p in pipeline] 27 | self.pipeline = Compose(pipeline) 28 | self.prefetch = prefetch 29 | 30 | def __len__(self): 31 | return self.data_source.get_length() 32 | 33 | @abstractmethod 34 | def __getitem__(self, idx): 35 | pass 36 | 37 | @abstractmethod 38 | def evaluate(self, scores, keyword, logger=None, **kwargs): 39 | pass 40 | -------------------------------------------------------------------------------- /openselfsup/datasets/builder.py: -------------------------------------------------------------------------------- 1 | import copy 2 | 3 | from openselfsup.utils import build_from_cfg 4 | from .dataset_wrappers import ConcatDataset, RepeatDataset 5 | from .registry import DATASETS, DATASOURCES 6 | 7 | 8 | def _concat_dataset(cfg, default_args=None): 9 | ann_files = cfg['ann_file'] 10 | img_prefixes = cfg.get('img_prefix', None) 11 | seg_prefixes = cfg.get('seg_prefix', None) 12 | proposal_files = cfg.get('proposal_file', None) 13 | 14 | datasets = [] 15 | num_dset = len(ann_files) 16 | for i in range(num_dset): 17 | data_cfg = copy.deepcopy(cfg) 18 | data_cfg['ann_file'] = ann_files[i] 19 | if isinstance(img_prefixes, (list, tuple)): 20 | data_cfg['img_prefix'] = img_prefixes[i] 21 | if isinstance(seg_prefixes, (list, tuple)): 22 | data_cfg['seg_prefix'] = seg_prefixes[i] 23 | if isinstance(proposal_files, (list, tuple)): 24 | data_cfg['proposal_file'] = proposal_files[i] 25 | datasets.append(build_dataset(data_cfg, default_args)) 26 | 27 | return ConcatDataset(datasets) 28 | 29 | 30 | def build_dataset(cfg, default_args=None): 31 | if isinstance(cfg, (list, tuple)): 32 | dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg]) 33 | elif cfg['type'] == 'RepeatDataset': 34 | dataset = RepeatDataset( 35 | build_dataset(cfg['dataset'], default_args), cfg['times']) 36 | else: 37 | dataset = build_from_cfg(cfg, DATASETS, default_args) 38 | 39 | return dataset 40 | 41 | 42 | def build_datasource(cfg): 43 | return build_from_cfg(cfg, DATASOURCES) 44 | -------------------------------------------------------------------------------- /openselfsup/datasets/byol.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import Dataset 3 | 4 | from openselfsup.utils import build_from_cfg 5 | 6 | from torchvision.transforms import Compose 7 | 8 | from .registry import DATASETS, PIPELINES 9 | from .builder import build_datasource 10 | from .utils import to_numpy 11 | 12 | 13 | @DATASETS.register_module 14 | class BYOLDataset(Dataset): 15 | """Dataset for BYOL. 16 | """ 17 | 18 | def __init__(self, data_source, pipeline1, pipeline2, prefetch=False): 19 | self.data_source = build_datasource(data_source) 20 | pipeline1 = [build_from_cfg(p, PIPELINES) for p in pipeline1] 21 | self.pipeline1 = Compose(pipeline1) 22 | pipeline2 = [build_from_cfg(p, PIPELINES) for p in pipeline2] 23 | self.pipeline2 = Compose(pipeline2) 24 | self.prefetch = prefetch 25 | 26 | def __len__(self): 27 | return self.data_source.get_length() 28 | 29 | def __getitem__(self, idx): 30 | img = self.data_source.get_sample(idx) 31 | img1 = self.pipeline1(img) 32 | img2 = self.pipeline2(img) 33 | if self.prefetch: 34 | img1 = torch.from_numpy(to_numpy(img1)) 35 | img2 = torch.from_numpy(to_numpy(img2)) 36 | 37 | img_cat = torch.cat((img1.unsqueeze(0), img2.unsqueeze(0)), dim=0) 38 | return dict(img=img_cat) 39 | 40 | def evaluate(self, scores, keyword, logger=None, **kwargs): 41 | raise NotImplemented 42 | -------------------------------------------------------------------------------- /openselfsup/datasets/classification.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from openselfsup.utils import print_log 4 | 5 | from .registry import DATASETS 6 | from .base import BaseDataset 7 | from .utils import to_numpy 8 | 9 | 10 | @DATASETS.register_module 11 | class ClassificationDataset(BaseDataset): 12 | """Dataset for classification. 13 | """ 14 | 15 | def __init__(self, data_source, pipeline, prefetch=False): 16 | super(ClassificationDataset, self).__init__(data_source, pipeline, prefetch) 17 | 18 | def __getitem__(self, idx): 19 | img, target = self.data_source.get_sample(idx) 20 | img = self.pipeline(img) 21 | if self.prefetch: 22 | img = torch.from_numpy(to_numpy(img)) 23 | return dict(img=img, gt_label=target) 24 | 25 | def evaluate(self, scores, keyword, logger=None, topk=(1, 5)): 26 | eval_res = {} 27 | 28 | target = torch.LongTensor(self.data_source.labels) 29 | assert scores.size(0) == target.size(0), \ 30 | "Inconsistent length for results and labels, {} vs {}".format( 31 | scores.size(0), target.size(0)) 32 | num = scores.size(0) 33 | _, pred = scores.topk(max(topk), dim=1, largest=True, sorted=True) 34 | pred = pred.t() 35 | correct = pred.eq(target.view(1, -1).expand_as(pred)) # KxN 36 | for k in topk: 37 | correct_k = correct[:k].view(-1).float().sum(0).item() 38 | acc = correct_k * 100.0 / num 39 | eval_res["{}_top{}".format(keyword, k)] = acc 40 | if logger is not None and logger != 'silent': 41 | print_log( 42 | "{}_top{}: {:.03f}".format(keyword, k, acc), 43 | logger=logger) 44 | return eval_res 45 | -------------------------------------------------------------------------------- /openselfsup/datasets/contrastive.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from PIL import Image 3 | from .registry import DATASETS 4 | from .base import BaseDataset 5 | from .utils import to_numpy 6 | 7 | 8 | @DATASETS.register_module 9 | class ContrastiveDataset(BaseDataset): 10 | """Dataset for contrastive learning methods that forward 11 | two views of the image at a time (MoCo, SimCLR). 12 | """ 13 | 14 | def __init__(self, data_source, pipeline, prefetch=False): 15 | data_source['return_label'] = False 16 | super(ContrastiveDataset, self).__init__(data_source, pipeline, prefetch) 17 | 18 | def __getitem__(self, idx): 19 | img = self.data_source.get_sample(idx) 20 | assert isinstance(img, Image.Image), \ 21 | 'The output from the data source must be an Image, got: {}. \ 22 | Please ensure that the list file does not contain labels.'.format( 23 | type(img)) 24 | img1 = self.pipeline(img) 25 | img2 = self.pipeline(img) 26 | if self.prefetch: 27 | img1 = torch.from_numpy(to_numpy(img1)) 28 | img2 = torch.from_numpy(to_numpy(img2)) 29 | img_cat = torch.cat((img1.unsqueeze(0), img2.unsqueeze(0)), dim=0) 30 | return dict(img=img_cat) 31 | 32 | def evaluate(self, scores, keyword, logger=None, **kwargs): 33 | raise NotImplemented 34 | -------------------------------------------------------------------------------- /openselfsup/datasets/data_sources/__init__.py: -------------------------------------------------------------------------------- 1 | from .cifar import Cifar10, Cifar100 2 | from .image_list import ImageList 3 | from .imagenet import ImageNet 4 | from .places205 import Places205 5 | -------------------------------------------------------------------------------- /openselfsup/datasets/data_sources/cifar.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | from PIL import Image 3 | 4 | from torchvision.datasets import CIFAR10, CIFAR100 5 | 6 | from ..registry import DATASOURCES 7 | 8 | 9 | class Cifar(metaclass=ABCMeta): 10 | 11 | CLASSES = None 12 | 13 | def __init__(self, root, split, return_label=True): 14 | assert split in ['train', 'test'] 15 | self.root = root 16 | self.split = split 17 | self.return_label = return_label 18 | self.cifar = None 19 | self.set_cifar() 20 | self.labels = self.cifar.targets 21 | 22 | @abstractmethod 23 | def set_cifar(self): 24 | pass 25 | 26 | def get_length(self): 27 | return len(self.cifar) 28 | 29 | def get_sample(self, idx): 30 | img = Image.fromarray(self.cifar.data[idx]) 31 | if self.return_label: 32 | target = self.labels[idx] # img: HWC, RGB 33 | return img, target 34 | else: 35 | return img 36 | 37 | 38 | @DATASOURCES.register_module 39 | class Cifar10(Cifar): 40 | 41 | CLASSES = [ 42 | 'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 43 | 'horse', 'ship', 'truck' 44 | ] 45 | 46 | def __init__(self, root, split, return_label=True): 47 | super().__init__(root, split, return_label) 48 | 49 | def set_cifar(self): 50 | try: 51 | self.cifar = CIFAR10( 52 | root=self.root, train=self.split == 'train', download=False) 53 | except: 54 | raise Exception("Please download CIFAR10 manually, \ 55 | in case of downloading the dataset parallelly \ 56 | that may corrupt the dataset.") 57 | 58 | 59 | @DATASOURCES.register_module 60 | class Cifar100(Cifar): 61 | 62 | def __init__(self, root, split, return_label=True): 63 | super().__init__(root, split, return_label) 64 | 65 | def set_cifar(self): 66 | try: 67 | self.cifar = CIFAR100( 68 | root=self.root, train=self.split == 'train', download=False) 69 | except: 70 | raise Exception("Please download CIFAR10 manually, \ 71 | in case of downloading the dataset parallelly \ 72 | that may corrupt the dataset.") 73 | -------------------------------------------------------------------------------- /openselfsup/datasets/data_sources/image_list.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | 4 | from ..registry import DATASOURCES 5 | from .utils import McLoader 6 | 7 | 8 | @DATASOURCES.register_module 9 | class ImageList(object): 10 | 11 | def __init__(self, root, list_file, memcached=False, mclient_path=None, return_label=True): 12 | with open(list_file, 'r') as f: 13 | lines = f.readlines() 14 | self.has_labels = len(lines[0].split()) == 2 15 | self.return_label = return_label 16 | if self.has_labels: 17 | self.fns, self.labels = zip(*[l.strip().split() for l in lines]) 18 | self.labels = [int(l) for l in self.labels] 19 | else: 20 | # assert self.return_label is False 21 | self.fns = [l.strip() for l in lines] 22 | self.fns = [os.path.join(root, fn) for fn in self.fns] 23 | self.memcached = memcached 24 | self.mclient_path = mclient_path 25 | self.initialized = False 26 | 27 | def _init_memcached(self): 28 | if not self.initialized: 29 | assert self.mclient_path is not None 30 | self.mc_loader = McLoader(self.mclient_path) 31 | self.initialized = True 32 | 33 | def get_length(self): 34 | return len(self.fns) 35 | 36 | def get_sample(self, idx): 37 | if self.memcached: 38 | self._init_memcached() 39 | if self.memcached: 40 | img = self.mc_loader(self.fns[idx]) 41 | else: 42 | img = Image.open(self.fns[idx]) 43 | img = img.convert('RGB') 44 | if self.has_labels and self.return_label: 45 | target = self.labels[idx] 46 | return img, target 47 | else: 48 | return img 49 | -------------------------------------------------------------------------------- /openselfsup/datasets/data_sources/imagenet.py: -------------------------------------------------------------------------------- 1 | from ..registry import DATASOURCES 2 | from .image_list import ImageList 3 | 4 | 5 | @DATASOURCES.register_module 6 | class ImageNet(ImageList): 7 | 8 | def __init__(self, root, list_file, memcached, mclient_path, return_label=True, *args, **kwargs): 9 | super(ImageNet, self).__init__( 10 | root, list_file, memcached, mclient_path, return_label) 11 | -------------------------------------------------------------------------------- /openselfsup/datasets/data_sources/places205.py: -------------------------------------------------------------------------------- 1 | from ..registry import DATASOURCES 2 | from .image_list import ImageList 3 | 4 | 5 | @DATASOURCES.register_module 6 | class Places205(ImageList): 7 | 8 | def __init__(self, root, list_file, memcached, mclient_path, return_label=True, *args, **kwargs): 9 | super(Places205, self).__init__( 10 | root, list_file, memcached, mclient_path, return_label) 11 | -------------------------------------------------------------------------------- /openselfsup/datasets/data_sources/utils.py: -------------------------------------------------------------------------------- 1 | import io 2 | from PIL import Image 3 | try: 4 | import mc 5 | except ImportError as E: 6 | pass 7 | 8 | 9 | def pil_loader(img_str): 10 | buff = io.BytesIO(img_str) 11 | return Image.open(buff) 12 | 13 | 14 | class McLoader(object): 15 | 16 | def __init__(self, mclient_path): 17 | assert mclient_path is not None, \ 18 | "Please specify 'data_mclient_path' in the config." 19 | self.mclient_path = mclient_path 20 | server_list_config_file = "{}/server_list.conf".format( 21 | self.mclient_path) 22 | client_config_file = "{}/client.conf".format(self.mclient_path) 23 | self.mclient = mc.MemcachedClient.GetInstance(server_list_config_file, 24 | client_config_file) 25 | 26 | def __call__(self, fn): 27 | try: 28 | img_value = mc.pyvector() 29 | self.mclient.Get(fn, img_value) 30 | img_value_str = mc.ConvertBuffer(img_value) 31 | img = pil_loader(img_value_str) 32 | except: 33 | print('Read image failed ({})'.format(fn)) 34 | return None 35 | else: 36 | return img 37 | -------------------------------------------------------------------------------- /openselfsup/datasets/dataset_wrappers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch.utils.data.dataset import ConcatDataset as _ConcatDataset 3 | 4 | from .registry import DATASETS 5 | 6 | 7 | @DATASETS.register_module 8 | class ConcatDataset(_ConcatDataset): 9 | """A wrapper of concatenated dataset. 10 | 11 | Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but 12 | concat the group flag for image aspect ratio. 13 | 14 | Args: 15 | datasets (list[:obj:`Dataset`]): A list of datasets. 16 | """ 17 | 18 | def __init__(self, datasets): 19 | super(ConcatDataset, self).__init__(datasets) 20 | self.CLASSES = datasets[0].CLASSES 21 | if hasattr(datasets[0], 'flag'): 22 | flags = [] 23 | for i in range(0, len(datasets)): 24 | flags.append(datasets[i].flag) 25 | self.flag = np.concatenate(flags) 26 | 27 | 28 | @DATASETS.register_module 29 | class RepeatDataset(object): 30 | """A wrapper of repeated dataset. 31 | 32 | The length of repeated dataset will be `times` larger than the original 33 | dataset. This is useful when the data loading time is long but the dataset 34 | is small. Using RepeatDataset can reduce the data loading time between 35 | epochs. 36 | 37 | Args: 38 | dataset (:obj:`Dataset`): The dataset to be repeated. 39 | times (int): Repeat times. 40 | """ 41 | 42 | def __init__(self, dataset, times): 43 | self.dataset = dataset 44 | self.times = times 45 | self.CLASSES = dataset.CLASSES 46 | if hasattr(self.dataset, 'flag'): 47 | self.flag = np.tile(self.dataset.flag, times) 48 | 49 | self._ori_len = len(self.dataset) 50 | 51 | def __getitem__(self, idx): 52 | return self.dataset[idx % self._ori_len] 53 | 54 | def __len__(self): 55 | return self.times * self._ori_len 56 | -------------------------------------------------------------------------------- /openselfsup/datasets/deepcluster.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | from .registry import DATASETS 3 | from .base import BaseDataset 4 | 5 | 6 | @DATASETS.register_module 7 | class DeepClusterDataset(BaseDataset): 8 | """Dataset for DC and ODC. 9 | """ 10 | 11 | def __init__(self, data_source, pipeline): 12 | super(DeepClusterDataset, self).__init__(data_source, pipeline) 13 | # init clustering labels 14 | self.labels = [-1 for _ in range(self.data_source.get_length())] 15 | 16 | def __getitem__(self, idx): 17 | img = self.data_source.get_sample(idx) 18 | assert isinstance(img, Image.Image), \ 19 | 'The output from the data source must be an Image, got: {}. \ 20 | Please ensure that the list file does not contain labels.'.format( 21 | type(img)) 22 | label = self.labels[idx] 23 | img = self.pipeline(img) 24 | return dict(img=img, pseudo_label=label, idx=idx) 25 | 26 | def assign_labels(self, labels): 27 | assert len(self.labels) == len(labels), \ 28 | "Inconsistent lenght of asigned labels, \ 29 | {} vs {}".format(len(self.labels), len(labels)) 30 | self.labels = labels[:] 31 | 32 | def evaluate(self, scores, keyword, logger=None): 33 | 34 | raise NotImplemented 35 | -------------------------------------------------------------------------------- /openselfsup/datasets/extraction.py: -------------------------------------------------------------------------------- 1 | from .registry import DATASETS 2 | from .base import BaseDataset 3 | 4 | 5 | @DATASETS.register_module 6 | class ExtractDataset(BaseDataset): 7 | """Dataset for feature extraction. 8 | """ 9 | 10 | def __init__(self, data_source, pipeline): 11 | super(ExtractDataset, self).__init__(data_source, pipeline) 12 | 13 | def __getitem__(self, idx): 14 | img = self.data_source.get_sample(idx) 15 | img = self.pipeline(img) 16 | return dict(img=img) 17 | 18 | def evaluate(self, scores, keyword, logger=None): 19 | raise NotImplemented 20 | -------------------------------------------------------------------------------- /openselfsup/datasets/loader/__init__.py: -------------------------------------------------------------------------------- 1 | from .build_loader import build_dataloader 2 | from .sampler import DistributedGroupSampler, GroupSampler, DistributedGivenIterationSampler 3 | 4 | __all__ = [ 5 | 'GroupSampler', 'DistributedGroupSampler', 'build_dataloader', 6 | 'DistributedGivenIterationSampler' 7 | ] 8 | -------------------------------------------------------------------------------- /openselfsup/datasets/npid.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | from .registry import DATASETS 3 | from .base import BaseDataset 4 | 5 | 6 | @DATASETS.register_module 7 | class NPIDDataset(BaseDataset): 8 | """Dataset for NPID. 9 | """ 10 | 11 | def __init__(self, data_source, pipeline): 12 | super(NPIDDataset, self).__init__(data_source, pipeline) 13 | 14 | def __getitem__(self, idx): 15 | img = self.data_source.get_sample(idx) 16 | assert isinstance(img, Image.Image), \ 17 | 'The output from the data source must be an Image, got: {}. \ 18 | Please ensure that the list file does not contain labels.'.format( 19 | type(img)) 20 | img = self.pipeline(img) 21 | return dict(img=img, idx=idx) 22 | 23 | def evaluate(self, scores, keyword, logger=None): 24 | 25 | raise NotImplemented 26 | -------------------------------------------------------------------------------- /openselfsup/datasets/pipelines/__init__.py: -------------------------------------------------------------------------------- 1 | from .transforms import * 2 | -------------------------------------------------------------------------------- /openselfsup/datasets/pipelines/transforms.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import inspect 3 | import numpy as np 4 | from PIL import Image, ImageFilter 5 | 6 | 7 | import torch 8 | from torchvision import transforms as _transforms 9 | 10 | from openselfsup.utils import build_from_cfg 11 | 12 | from ..registry import PIPELINES 13 | 14 | # register all existing transforms in torchvision 15 | _EXCLUDED_TRANSFORMS = ['GaussianBlur'] 16 | for m in inspect.getmembers(_transforms, inspect.isclass): 17 | if m[0] not in _EXCLUDED_TRANSFORMS: 18 | PIPELINES.register_module(m[1]) 19 | 20 | 21 | @PIPELINES.register_module 22 | class RandomAppliedTrans(object): 23 | """Randomly applied transformations. 24 | 25 | Args: 26 | transforms (list[dict]): List of transformations in dictionaries. 27 | p (float): Probability. 28 | """ 29 | 30 | def __init__(self, transforms, p=0.5): 31 | t = [build_from_cfg(t, PIPELINES) for t in transforms] 32 | self.trans = _transforms.RandomApply(t, p=p) 33 | 34 | def __call__(self, img): 35 | return self.trans(img) 36 | 37 | def __repr__(self): 38 | repr_str = self.__class__.__name__ 39 | return repr_str 40 | 41 | 42 | # custom transforms 43 | @PIPELINES.register_module 44 | class Lighting(object): 45 | """Lighting noise(AlexNet - style PCA - based noise).""" 46 | 47 | _IMAGENET_PCA = { 48 | 'eigval': 49 | torch.Tensor([0.2175, 0.0188, 0.0045]), 50 | 'eigvec': 51 | torch.Tensor([ 52 | [-0.5675, 0.7192, 0.4009], 53 | [-0.5808, -0.0045, -0.8140], 54 | [-0.5836, -0.6948, 0.4203], 55 | ]) 56 | } 57 | 58 | def __init__(self): 59 | self.alphastd = 0.1 60 | self.eigval = self._IMAGENET_PCA['eigval'] 61 | self.eigvec = self._IMAGENET_PCA['eigvec'] 62 | 63 | def __call__(self, img): 64 | assert isinstance(img, torch.Tensor), \ 65 | "Expect torch.Tensor, got {}".format(type(img)) 66 | if self.alphastd == 0: 67 | return img 68 | 69 | alpha = img.new().resize_(3).normal_(0, self.alphastd) 70 | rgb = self.eigvec.type_as(img).clone()\ 71 | .mul(alpha.view(1, 3).expand(3, 3))\ 72 | .mul(self.eigval.view(1, 3).expand(3, 3))\ 73 | .sum(1).squeeze() 74 | 75 | return img.add(rgb.view(3, 1, 1).expand_as(img)) 76 | 77 | def __repr__(self): 78 | repr_str = self.__class__.__name__ 79 | return repr_str 80 | 81 | 82 | @PIPELINES.register_module 83 | class GaussianBlur(object): 84 | """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709.""" 85 | 86 | def __init__(self, sigma_min, sigma_max): 87 | self.sigma_min = sigma_min 88 | self.sigma_max = sigma_max 89 | 90 | def __call__(self, img): 91 | sigma = np.random.uniform(self.sigma_min, self.sigma_max) 92 | img = img.filter(ImageFilter.GaussianBlur(radius=sigma)) 93 | return img 94 | 95 | def __repr__(self): 96 | repr_str = self.__class__.__name__ 97 | return repr_str 98 | 99 | 100 | @PIPELINES.register_module 101 | class Solarization(object): 102 | """Solarization augmentation in BYOL https://arxiv.org/abs/2006.07733.""" 103 | 104 | def __init__(self, threshold=128): 105 | self.threshold = threshold 106 | 107 | def __call__(self, img): 108 | img = np.array(img) 109 | img = np.where(img < self.threshold, img, 255 -img) 110 | return Image.fromarray(img.astype(np.uint8)) 111 | 112 | def __repr__(self): 113 | repr_str = self.__class__.__name__ 114 | return repr_str 115 | -------------------------------------------------------------------------------- /openselfsup/datasets/registry.py: -------------------------------------------------------------------------------- 1 | from openselfsup.utils import Registry 2 | 3 | DATASOURCES = Registry('datasource') 4 | DATASETS = Registry('dataset') 5 | PIPELINES = Registry('pipeline') 6 | -------------------------------------------------------------------------------- /openselfsup/datasets/relative_loc.py: -------------------------------------------------------------------------------- 1 | from openselfsup.utils import build_from_cfg 2 | 3 | import torch 4 | from PIL import Image 5 | from torchvision.transforms import Compose, RandomCrop 6 | import torchvision.transforms.functional as TF 7 | 8 | from .registry import DATASETS, PIPELINES 9 | from .base import BaseDataset 10 | 11 | 12 | def image_to_patches(img): 13 | """Crop split_per_side x split_per_side patches from input image. 14 | 15 | Args: 16 | img (PIL Image): input image. 17 | 18 | Returns: 19 | list[PIL Image]: A list of cropped patches. 20 | """ 21 | split_per_side = 3 # split of patches per image side 22 | patch_jitter = 21 # jitter of each patch from each grid 23 | h, w = img.size 24 | h_grid = h // split_per_side 25 | w_grid = w // split_per_side 26 | h_patch = h_grid - patch_jitter 27 | w_patch = w_grid - patch_jitter 28 | assert h_patch > 0 and w_patch > 0 29 | patches = [] 30 | for i in range(split_per_side): 31 | for j in range(split_per_side): 32 | p = TF.crop(img, i * h_grid, j * w_grid, h_grid, w_grid) 33 | p = RandomCrop((h_patch, w_patch))(p) 34 | patches.append(p) 35 | return patches 36 | 37 | 38 | @DATASETS.register_module 39 | class RelativeLocDataset(BaseDataset): 40 | """Dataset for relative patch location. 41 | """ 42 | 43 | def __init__(self, data_source, pipeline, format_pipeline): 44 | super(RelativeLocDataset, self).__init__(data_source, pipeline) 45 | format_pipeline = [build_from_cfg(p, PIPELINES) for p in format_pipeline] 46 | self.format_pipeline = Compose(format_pipeline) 47 | 48 | def __getitem__(self, idx): 49 | img = self.data_source.get_sample(idx) 50 | assert isinstance(img, Image.Image), \ 51 | 'The output from the data source must be an Image, got: {}. \ 52 | Please ensure that the list file does not contain labels.'.format( 53 | type(img)) 54 | img = self.pipeline(img) 55 | patches = image_to_patches(img) 56 | patches = [self.format_pipeline(p) for p in patches] 57 | perms = [] 58 | # create a list of patch pairs 59 | [perms.append(torch.cat((patches[i], patches[4]), dim=0)) for i in range(9) if i != 4] 60 | # create corresponding labels for patch pairs 61 | patch_labels = torch.LongTensor([0, 1, 2, 3, 4, 5, 6, 7]) 62 | return dict(img=torch.stack(perms), patch_label=patch_labels) # 8(2C)HW, 8 63 | 64 | def evaluate(self, scores, keyword, logger=None): 65 | raise NotImplemented 66 | -------------------------------------------------------------------------------- /openselfsup/datasets/rotation_pred.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from PIL import Image 3 | 4 | from .registry import DATASETS 5 | from .base import BaseDataset 6 | 7 | 8 | def rotate(img): 9 | """Rotate input image with 0, 90, 180, and 270 degrees. 10 | 11 | Args: 12 | img (Tensor): input image of shape (C, H, W). 13 | 14 | Returns: 15 | list[Tensor]: A list of four rotated images. 16 | """ 17 | return [ 18 | img, 19 | torch.flip(img.transpose(1, 2), [1]), 20 | torch.flip(img, [1, 2]), 21 | torch.flip(img, [1]).transpose(1, 2) 22 | ] 23 | 24 | 25 | @DATASETS.register_module 26 | class RotationPredDataset(BaseDataset): 27 | """Dataset for rotation prediction. 28 | """ 29 | 30 | def __init__(self, data_source, pipeline): 31 | super(RotationPredDataset, self).__init__(data_source, pipeline) 32 | 33 | def __getitem__(self, idx): 34 | img = self.data_source.get_sample(idx) 35 | assert isinstance(img, Image.Image), \ 36 | 'The output from the data source must be an Image, got: {}. \ 37 | Please ensure that the list file does not contain labels.'.format( 38 | type(img)) 39 | img = self.pipeline(img) 40 | img = torch.stack(rotate(img), dim=0) 41 | rotation_labels = torch.LongTensor([0, 1, 2, 3]) 42 | return dict(img=img, rot_label=rotation_labels) 43 | 44 | def evaluate(self, scores, keyword, logger=None): 45 | raise NotImplemented 46 | -------------------------------------------------------------------------------- /openselfsup/datasets/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def to_numpy(pil_img): 5 | np_img = np.array(pil_img, dtype=np.uint8) 6 | if np_img.ndim < 3: 7 | np_img = np.expand_dims(np_img, axis=-1) 8 | np_img = np.rollaxis(np_img, 2) # HWC to CHW 9 | return np_img 10 | -------------------------------------------------------------------------------- /openselfsup/hooks/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_hook 2 | from .byol_hook import BYOLHook 3 | from .deepcluster_hook import DeepClusterHook 4 | from .odc_hook import ODCHook 5 | from .optimizer_hook import DistOptimizerHook 6 | from .extractor import Extractor 7 | from .validate_hook import ValidateHook 8 | from .registry import HOOKS 9 | -------------------------------------------------------------------------------- /openselfsup/hooks/builder.py: -------------------------------------------------------------------------------- 1 | from openselfsup.utils import build_from_cfg 2 | 3 | from .registry import HOOKS 4 | 5 | 6 | def build_hook(cfg, default_args=None): 7 | return build_from_cfg(cfg, HOOKS, default_args) 8 | -------------------------------------------------------------------------------- /openselfsup/hooks/byol_hook.py: -------------------------------------------------------------------------------- 1 | from math import cos, pi 2 | from mmcv.runner import Hook 3 | from mmcv.parallel import is_module_wrapper 4 | 5 | from .registry import HOOKS 6 | 7 | 8 | @HOOKS.register_module 9 | class BYOLHook(Hook): 10 | """Hook for BYOL. 11 | 12 | This hook includes momentum adjustment in BYOL following: 13 | m = 1 - ( 1- m_0) * (cos(pi * k / K) + 1) / 2, 14 | k: current step, K: total steps. 15 | 16 | Args: 17 | end_momentum (float): The final momentum coefficient 18 | for the target network. Default: 1. 19 | """ 20 | 21 | def __init__(self, end_momentum=1., update_interval=1, **kwargs): 22 | self.end_momentum = end_momentum 23 | self.update_interval = update_interval 24 | 25 | def before_train_iter(self, runner): 26 | assert hasattr(runner.model.module, 'momentum'), \ 27 | "The runner must have attribute \"momentum\" in BYOLHook." 28 | assert hasattr(runner.model.module, 'base_momentum'), \ 29 | "The runner must have attribute \"base_momentum\" in BYOLHook." 30 | if self.every_n_iters(runner, self.update_interval): 31 | cur_iter = runner.iter 32 | max_iter = runner.max_iters 33 | base_m = runner.model.module.base_momentum 34 | m = self.end_momentum - (self.end_momentum - base_m) * ( 35 | cos(pi * cur_iter / float(max_iter)) + 1) / 2 36 | runner.model.module.momentum = m 37 | 38 | def after_train_iter(self, runner): 39 | if self.every_n_iters(runner, self.update_interval): 40 | if is_module_wrapper(runner.model): 41 | runner.model.module.momentum_update() 42 | else: 43 | runner.model.momentum_update() 44 | -------------------------------------------------------------------------------- /openselfsup/hooks/extractor.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch.utils.data import Dataset 3 | 4 | from openselfsup.utils import nondist_forward_collect, dist_forward_collect 5 | 6 | 7 | class Extractor(object): 8 | """Feature extractor. 9 | 10 | Args: 11 | dataset (Dataset | dict): A PyTorch dataset or dict that indicates 12 | the dataset. 13 | imgs_per_gpu (int): Number of images on each GPU, i.e., batch size of 14 | each GPU. 15 | workers_per_gpu (int): How many subprocesses to use for data loading 16 | for each GPU. 17 | dist_mode (bool): Use distributed extraction or not. Default: False. 18 | """ 19 | 20 | def __init__(self, 21 | dataset, 22 | imgs_per_gpu, 23 | workers_per_gpu, 24 | dist_mode=False): 25 | from openselfsup import datasets 26 | if isinstance(dataset, Dataset): 27 | self.dataset = dataset 28 | elif isinstance(dataset, dict): 29 | self.dataset = datasets.build_dataset(dataset) 30 | else: 31 | raise TypeError( 32 | 'dataset must be a Dataset object or a dict, not {}'.format( 33 | type(dataset))) 34 | self.data_loader = datasets.build_dataloader( 35 | self.dataset, 36 | imgs_per_gpu, 37 | workers_per_gpu, 38 | dist=dist_mode, 39 | shuffle=False) 40 | self.dist_mode = dist_mode 41 | self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) 42 | 43 | def _forward_func(self, runner, **x): 44 | backbone_feat = runner.model(mode='extract', **x) 45 | last_layer_feat = runner.model.module.neck([backbone_feat[-1]])[0] 46 | last_layer_feat = last_layer_feat.view(last_layer_feat.size(0), -1) 47 | return dict(feature=last_layer_feat.cpu()) 48 | 49 | def __call__(self, runner): 50 | func = lambda **x: self._forward_func(runner, **x) 51 | if self.dist_mode: 52 | feats = dist_forward_collect( 53 | func, 54 | self.data_loader, 55 | runner.rank, 56 | len(self.dataset), 57 | ret_rank=-1)['feature'] # NxD 58 | else: 59 | feats = nondist_forward_collect(func, self.data_loader, 60 | len(self.dataset))['feature'] 61 | return feats 62 | -------------------------------------------------------------------------------- /openselfsup/hooks/odc_hook.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from mmcv.runner import Hook 4 | 5 | from openselfsup.utils import print_log 6 | from .registry import HOOKS 7 | 8 | 9 | @HOOKS.register_module 10 | class ODCHook(Hook): 11 | """Hook for ODC. 12 | 13 | Args: 14 | centroids_update_interval (int): Frequency of iterations 15 | to update centroids. 16 | deal_with_small_clusters_interval (int): Frequency of iterations 17 | to deal with small clusters. 18 | evaluate_interval (int): Frequency of iterations to evaluate clusters. 19 | reweight (bool): Whether to perform loss re-weighting. 20 | reweight_pow (float): The power of re-weighting. 21 | dist_mode (bool): Use distributed training or not. Default: True. 22 | """ 23 | 24 | def __init__(self, 25 | centroids_update_interval, 26 | deal_with_small_clusters_interval, 27 | evaluate_interval, 28 | reweight, 29 | reweight_pow, 30 | dist_mode=True): 31 | assert dist_mode, "non-dist mode is not implemented" 32 | self.centroids_update_interval = centroids_update_interval 33 | self.deal_with_small_clusters_interval = \ 34 | deal_with_small_clusters_interval 35 | self.evaluate_interval = evaluate_interval 36 | self.reweight = reweight 37 | self.reweight_pow = reweight_pow 38 | 39 | def after_train_iter(self, runner): 40 | # centroids update 41 | if self.every_n_iters(runner, self.centroids_update_interval): 42 | runner.model.module.memory_bank.update_centroids_memory() 43 | 44 | # deal with small clusters 45 | if self.every_n_iters(runner, self.deal_with_small_clusters_interval): 46 | runner.model.module.memory_bank.deal_with_small_clusters() 47 | 48 | # reweight 49 | runner.model.module.set_reweight() 50 | 51 | # evaluate 52 | if self.every_n_iters(runner, self.evaluate_interval): 53 | new_labels = runner.model.module.memory_bank.label_bank 54 | if new_labels.is_cuda: 55 | new_labels = new_labels.cpu() 56 | self.evaluate(runner, new_labels.numpy()) 57 | 58 | def after_train_epoch(self, runner): 59 | # save cluster 60 | if self.every_n_epochs(runner, 10) and runner.rank == 0: 61 | new_labels = runner.model.module.memory_bank.label_bank 62 | if new_labels.is_cuda: 63 | new_labels = new_labels.cpu() 64 | np.save( 65 | "{}/cluster_epoch_{}.npy".format(runner.work_dir, 66 | runner.epoch), 67 | new_labels.numpy()) 68 | 69 | def evaluate(self, runner, new_labels): 70 | hist = np.bincount( 71 | new_labels, minlength=runner.model.module.memory_bank.num_classes) 72 | empty_cls = (hist == 0).sum() 73 | minimal_cls_size, maximal_cls_size = hist.min(), hist.max() 74 | if runner.rank == 0: 75 | print_log( 76 | "empty_num: {}\tmin_cluster: {}\tmax_cluster:{}".format( 77 | empty_cls.item(), minimal_cls_size.item(), 78 | maximal_cls_size.item()), 79 | logger='root') 80 | -------------------------------------------------------------------------------- /openselfsup/hooks/optimizer_hook.py: -------------------------------------------------------------------------------- 1 | from mmcv.runner import OptimizerHook 2 | try: 3 | import apex 4 | except: 5 | print('apex is not installed') 6 | 7 | class DistOptimizerHook(OptimizerHook): 8 | """Optimizer hook for distributed training.""" 9 | 10 | def __init__(self, update_interval=1, grad_clip=None, coalesce=True, bucket_size_mb=-1, use_fp16=False): 11 | self.grad_clip = grad_clip 12 | self.coalesce = coalesce 13 | self.bucket_size_mb = bucket_size_mb 14 | self.update_interval = update_interval 15 | self.use_fp16 = use_fp16 16 | 17 | def before_run(self, runner): 18 | runner.optimizer.zero_grad() 19 | 20 | def after_train_iter(self, runner): 21 | runner.outputs['loss'] /= self.update_interval 22 | if self.use_fp16: 23 | with apex.amp.scale_loss(runner.outputs['loss'], runner.optimizer) as scaled_loss: 24 | scaled_loss.backward() 25 | else: 26 | runner.outputs['loss'].backward() 27 | if self.every_n_iters(runner, self.update_interval): 28 | if self.grad_clip is not None: 29 | self.clip_grads(runner.model.parameters()) 30 | runner.optimizer.step() 31 | runner.optimizer.zero_grad() 32 | -------------------------------------------------------------------------------- /openselfsup/hooks/registry.py: -------------------------------------------------------------------------------- 1 | from openselfsup.utils import Registry 2 | 3 | HOOKS = Registry('hook') 4 | -------------------------------------------------------------------------------- /openselfsup/hooks/validate_hook.py: -------------------------------------------------------------------------------- 1 | from mmcv.runner import Hook 2 | 3 | import torch 4 | from torch.utils.data import Dataset 5 | 6 | from openselfsup.utils import nondist_forward_collect, dist_forward_collect 7 | from .registry import HOOKS 8 | 9 | 10 | @HOOKS.register_module 11 | class ValidateHook(Hook): 12 | """Validation hook. 13 | 14 | Args: 15 | dataset (Dataset | dict): A PyTorch dataset or dict that indicates 16 | the dataset. 17 | dist_mode (bool): Use distributed evaluation or not. Default: True. 18 | initial (bool): Whether to evaluate before the training starts. 19 | Default: True. 20 | interval (int): Evaluation interval (by epochs). Default: 1. 21 | **eval_kwargs: Evaluation arguments fed into the evaluate function of 22 | the dataset. 23 | """ 24 | 25 | def __init__(self, 26 | dataset, 27 | dist_mode=True, 28 | initial=True, 29 | interval=1, 30 | **eval_kwargs): 31 | from openselfsup import datasets 32 | if isinstance(dataset, Dataset): 33 | self.dataset = dataset 34 | elif isinstance(dataset, dict): 35 | self.dataset = datasets.build_dataset(dataset) 36 | else: 37 | raise TypeError( 38 | 'dataset must be a Dataset object or a dict, not {}'.format( 39 | type(dataset))) 40 | self.data_loader = datasets.build_dataloader( 41 | self.dataset, 42 | eval_kwargs['imgs_per_gpu'], 43 | eval_kwargs['workers_per_gpu'], 44 | dist=dist_mode, 45 | shuffle=False, 46 | prefetch=eval_kwargs.get('prefetch', False), 47 | img_norm_cfg=eval_kwargs.get('img_norm_cfg', dict()), 48 | ) 49 | self.dist_mode = dist_mode 50 | self.initial = initial 51 | self.interval = interval 52 | self.eval_kwargs = eval_kwargs 53 | 54 | def before_run(self, runner): 55 | if self.initial: 56 | self._run_validate(runner) 57 | 58 | def after_train_epoch(self, runner): 59 | if not self.every_n_epochs(runner, self.interval): 60 | return 61 | self._run_validate(runner) 62 | 63 | def _run_validate(self, runner): 64 | runner.model.eval() 65 | func = lambda **x: runner.model(mode='test', **x) 66 | if self.dist_mode: 67 | results = dist_forward_collect( 68 | func, self.data_loader, runner.rank, 69 | len(self.dataset)) # dict{key: np.ndarray} 70 | else: 71 | results = nondist_forward_collect(func, self.data_loader, 72 | len(self.dataset)) 73 | if runner.rank == 0: 74 | for name, val in results.items(): 75 | self._evaluate(runner, torch.from_numpy(val), name) 76 | runner.model.train() 77 | 78 | def _evaluate(self, runner, results, keyword): 79 | eval_res = self.dataset.evaluate( 80 | results, 81 | keyword=keyword, 82 | logger=runner.logger, 83 | **self.eval_kwargs['eval_param']) 84 | for name, val in eval_res.items(): 85 | runner.log_buffer.output[name] = val 86 | runner.log_buffer.ready = True 87 | -------------------------------------------------------------------------------- /openselfsup/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .backbones import * # noqa: F401,F403 2 | from .builder import (build_backbone, build_model, build_head, build_loss) 3 | from .byol import BYOL 4 | from .heads import * 5 | from .classification import Classification 6 | from .deepcluster import DeepCluster 7 | from .odc import ODC 8 | from .necks import * 9 | from .npid import NPID 10 | from .memories import * 11 | from .moco import MOCO 12 | from .RegionCLM import RegionCLM 13 | from .registry import (BACKBONES, MODELS, NECKS, MEMORIES, HEADS, LOSSES) 14 | from .rotation_pred import RotationPred 15 | from .relative_loc import RelativeLoc 16 | from .simclr import SimCLR 17 | -------------------------------------------------------------------------------- /openselfsup/models/backbones/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import ResNet, make_res_layer 2 | -------------------------------------------------------------------------------- /openselfsup/models/builder.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from openselfsup.utils import build_from_cfg 4 | from .registry import (BACKBONES, MODELS, NECKS, HEADS, MEMORIES, LOSSES) 5 | 6 | 7 | def build(cfg, registry, default_args=None): 8 | """Build a module. 9 | 10 | Args: 11 | cfg (dict, list[dict]): The config of modules, it is either a dict 12 | or a list of configs. 13 | registry (:obj:`Registry`): A registry the module belongs to. 14 | default_args (dict, optional): Default arguments to build the module. 15 | Default: None. 16 | 17 | Returns: 18 | nn.Module: A built nn module. 19 | """ 20 | if isinstance(cfg, list): 21 | modules = [ 22 | build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg 23 | ] 24 | return nn.Sequential(*modules) 25 | else: 26 | return build_from_cfg(cfg, registry, default_args) 27 | 28 | 29 | def build_backbone(cfg): 30 | """Build backbone.""" 31 | return build(cfg, BACKBONES) 32 | 33 | 34 | def build_neck(cfg): 35 | """Build neck.""" 36 | return build(cfg, NECKS) 37 | 38 | 39 | def build_memory(cfg): 40 | """Build memory.""" 41 | return build(cfg, MEMORIES) 42 | 43 | 44 | def build_head(cfg): 45 | """Build head.""" 46 | return build(cfg, HEADS) 47 | 48 | 49 | def build_loss(cfg): 50 | """Build loss.""" 51 | return build(cfg, LOSSES) 52 | 53 | 54 | def build_model(cfg): 55 | """Build model.""" 56 | return build(cfg, MODELS) 57 | -------------------------------------------------------------------------------- /openselfsup/models/heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .contrastive_head import ContrastiveHead 2 | from .cls_head import ClsHead 3 | from .latent_pred_head import LatentPredictHead 4 | from .multi_cls_head import MultiClsHead 5 | -------------------------------------------------------------------------------- /openselfsup/models/heads/cls_head.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from mmcv.cnn import kaiming_init, normal_init 3 | 4 | from ..utils import accuracy 5 | from ..registry import HEADS 6 | 7 | 8 | @HEADS.register_module 9 | class ClsHead(nn.Module): 10 | """Simplest classifier head, with only one fc layer. 11 | """ 12 | 13 | def __init__(self, 14 | with_avg_pool=False, 15 | in_channels=2048, 16 | num_classes=1000): 17 | super(ClsHead, self).__init__() 18 | self.with_avg_pool = with_avg_pool 19 | self.in_channels = in_channels 20 | self.num_classes = num_classes 21 | 22 | self.criterion = nn.CrossEntropyLoss() 23 | 24 | if self.with_avg_pool: 25 | self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) 26 | self.fc_cls = nn.Linear(in_channels, num_classes) 27 | 28 | def init_weights(self, init_linear='normal', std=0.01, bias=0.): 29 | assert init_linear in ['normal', 'kaiming'], \ 30 | "Undefined init_linear: {}".format(init_linear) 31 | for m in self.modules(): 32 | if isinstance(m, nn.Linear): 33 | if init_linear == 'normal': 34 | normal_init(m, std=std, bias=bias) 35 | else: 36 | kaiming_init(m, mode='fan_in', nonlinearity='relu') 37 | elif isinstance(m, 38 | (nn.BatchNorm2d, nn.GroupNorm, nn.SyncBatchNorm)): 39 | if m.weight is not None: 40 | nn.init.constant_(m.weight, 1) 41 | if m.bias is not None: 42 | nn.init.constant_(m.bias, 0) 43 | 44 | def forward(self, x): 45 | assert isinstance(x, (tuple, list)) and len(x) == 1 46 | x = x[0] 47 | if self.with_avg_pool: 48 | assert x.dim() == 4, \ 49 | "Tensor must has 4 dims, got: {}".format(x.dim()) 50 | x = self.avg_pool(x) 51 | x = x.view(x.size(0), -1) 52 | cls_score = self.fc_cls(x) 53 | return [cls_score] 54 | 55 | def loss(self, cls_score, labels): 56 | losses = dict() 57 | assert isinstance(cls_score, (tuple, list)) and len(cls_score) == 1 58 | losses['loss'] = self.criterion(cls_score[0], labels) 59 | losses['acc'] = accuracy(cls_score[0], labels) 60 | return losses 61 | -------------------------------------------------------------------------------- /openselfsup/models/heads/contrastive_head.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from ..registry import HEADS 5 | 6 | 7 | @HEADS.register_module 8 | class ContrastiveHead(nn.Module): 9 | """Head for contrastive learning. 10 | 11 | Args: 12 | temperature (float): The temperature hyper-parameter that 13 | controls the concentration level of the distribution. 14 | Default: 0.1. 15 | """ 16 | 17 | def __init__(self, temperature=0.1): 18 | super(ContrastiveHead, self).__init__() 19 | self.criterion = nn.CrossEntropyLoss() 20 | self.temperature = temperature 21 | 22 | def forward(self, pos, neg): 23 | """Forward head. 24 | 25 | Args: 26 | pos (Tensor): Nx1 positive similarity. 27 | neg (Tensor): Nxk negative similarity. 28 | 29 | Returns: 30 | dict[str, Tensor]: A dictionary of loss components. 31 | """ 32 | N = pos.size(0) 33 | logits = torch.cat((pos, neg), dim=1) 34 | logits /= self.temperature 35 | labels = torch.zeros((N, ), dtype=torch.long).cuda() 36 | losses = dict() 37 | losses['loss'] = self.criterion(logits, labels) 38 | return losses 39 | -------------------------------------------------------------------------------- /openselfsup/models/heads/latent_pred_head.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from mmcv.cnn import normal_init 4 | 5 | from ..registry import HEADS 6 | from .. import builder 7 | 8 | @HEADS.register_module 9 | class LatentPredictHead(nn.Module): 10 | """Head for contrastive learning. 11 | """ 12 | 13 | def __init__(self, predictor, size_average=True): 14 | super(LatentPredictHead, self).__init__() 15 | self.predictor = builder.build_neck(predictor) 16 | self.size_average = size_average 17 | 18 | def init_weights(self, init_linear='normal'): 19 | self.predictor.init_weights(init_linear=init_linear) 20 | 21 | def forward(self, input, target): 22 | """Forward head. 23 | 24 | Args: 25 | input (Tensor): NxC input features. 26 | target (Tensor): NxC target features. 27 | 28 | Returns: 29 | dict[str, Tensor]: A dictionary of loss components. 30 | """ 31 | pred = self.predictor([input])[0] 32 | pred_norm = nn.functional.normalize(pred, dim=1) 33 | target_norm = nn.functional.normalize(target, dim=1) 34 | loss = -2 * (pred_norm * target_norm).sum() 35 | if self.size_average: 36 | loss /= input.size(0) 37 | return dict(loss=loss) 38 | 39 | 40 | @HEADS.register_module 41 | class LatentClsHead(nn.Module): 42 | """Head for contrastive learning. 43 | """ 44 | 45 | def __init__(self, predictor): 46 | super(LatentClsHead, self).__init__() 47 | self.predictor = nn.Linear(predictor.in_channels, 48 | predictor.num_classes) 49 | self.criterion = nn.CrossEntropyLoss() 50 | 51 | def init_weights(self, init_linear='normal'): 52 | normal_init(self.predictor, std=0.01) 53 | 54 | def forward(self, input, target): 55 | """Forward head. 56 | 57 | Args: 58 | input (Tensor): NxC input features. 59 | target (Tensor): NxC target features. 60 | 61 | Returns: 62 | dict[str, Tensor]: A dictionary of loss components. 63 | """ 64 | pred = self.predictor(input) 65 | with torch.no_grad(): 66 | label = torch.argmax(self.predictor(target), dim=1).detach() 67 | loss = self.criterion(pred, label) 68 | return dict(loss=loss) 69 | -------------------------------------------------------------------------------- /openselfsup/models/heads/multi_cls_head.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | from ..utils import accuracy 4 | from ..registry import HEADS 5 | from ..utils import build_norm_layer, MultiPooling 6 | 7 | 8 | @HEADS.register_module 9 | class MultiClsHead(nn.Module): 10 | """Multiple classifier heads. 11 | """ 12 | 13 | FEAT_CHANNELS = {'resnet50': [64, 256, 512, 1024, 2048]} 14 | FEAT_LAST_UNPOOL = {'resnet50': 2048 * 7 * 7} 15 | 16 | def __init__(self, 17 | pool_type='adaptive', 18 | in_indices=(0, ), 19 | with_last_layer_unpool=False, 20 | backbone='resnet50', 21 | norm_cfg=dict(type='BN'), 22 | num_classes=1000): 23 | super(MultiClsHead, self).__init__() 24 | assert norm_cfg['type'] in ['BN', 'SyncBN', 'GN', 'null'] 25 | 26 | self.with_last_layer_unpool = with_last_layer_unpool 27 | self.with_norm = norm_cfg['type'] != 'null' 28 | 29 | self.criterion = nn.CrossEntropyLoss() 30 | 31 | self.multi_pooling = MultiPooling(pool_type, in_indices, backbone) 32 | 33 | if self.with_norm: 34 | self.norms = nn.ModuleList([ 35 | build_norm_layer(norm_cfg, self.FEAT_CHANNELS[backbone][l])[1] 36 | for l in in_indices 37 | ]) 38 | 39 | self.fcs = nn.ModuleList([ 40 | nn.Linear(self.multi_pooling.POOL_DIMS[backbone][l], num_classes) 41 | for l in in_indices 42 | ]) 43 | if with_last_layer_unpool: 44 | self.fcs.append( 45 | nn.Linear(self.FEAT_LAST_UNPOOL[backbone], num_classes)) 46 | 47 | def init_weights(self): 48 | for m in self.modules(): 49 | if isinstance(m, nn.Linear): 50 | nn.init.normal_(m.weight, 0, 0.01) 51 | nn.init.constant_(m.bias, 0) 52 | elif isinstance(m, 53 | (nn.BatchNorm2d, nn.GroupNorm, nn.SyncBatchNorm)): 54 | if m.weight is not None: 55 | nn.init.constant_(m.weight, 1) 56 | if m.bias is not None: 57 | nn.init.constant_(m.bias, 0) 58 | 59 | def forward(self, x): 60 | assert isinstance(x, (list, tuple)) 61 | if self.with_last_layer_unpool: 62 | last_x = x[-1] 63 | x = self.multi_pooling(x) 64 | if self.with_norm: 65 | x = [n(xx) for n, xx in zip(self.norms, x)] 66 | if self.with_last_layer_unpool: 67 | x.append(last_x) 68 | x = [xx.view(xx.size(0), -1) for xx in x] 69 | x = [fc(xx) for fc, xx in zip(self.fcs, x)] 70 | return x 71 | 72 | def loss(self, cls_score, labels): 73 | losses = dict() 74 | for i, s in enumerate(cls_score): 75 | # keys must contain "loss" 76 | losses['loss.{}'.format(i + 1)] = self.criterion(s, labels) 77 | losses['acc.{}'.format(i + 1)] = accuracy(s, labels) 78 | return losses 79 | -------------------------------------------------------------------------------- /openselfsup/models/memories/__init__.py: -------------------------------------------------------------------------------- 1 | from .odc_memory import ODCMemory 2 | from .simple_memory import SimpleMemory 3 | -------------------------------------------------------------------------------- /openselfsup/models/memories/simple_memory.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.distributed as dist 4 | from mmcv.runner import get_dist_info 5 | from openselfsup.utils import AliasMethod 6 | 7 | from ..registry import MEMORIES 8 | 9 | 10 | @MEMORIES.register_module 11 | class SimpleMemory(nn.Module): 12 | """Simple memory bank for NPID. 13 | 14 | Args: 15 | length (int): Number of features stored in the memory bank. 16 | feat_dim (int): Dimension of stored features. 17 | momentum (float): Momentum coefficient for updating features. 18 | """ 19 | 20 | def __init__(self, length, feat_dim, momentum, **kwargs): 21 | super(SimpleMemory, self).__init__() 22 | self.rank, self.num_replicas = get_dist_info() 23 | self.feature_bank = torch.randn(length, feat_dim).cuda() 24 | self.feature_bank = nn.functional.normalize(self.feature_bank) 25 | self.momentum = momentum 26 | self.multinomial = AliasMethod(torch.ones(length)) 27 | self.multinomial.cuda() 28 | 29 | def update(self, ind, feature): 30 | """Update features in memory bank. 31 | 32 | Args: 33 | ind (Tensor): Indices for the batch of features. 34 | feature (Tensor): Batch of features. 35 | """ 36 | feature_norm = nn.functional.normalize(feature) 37 | ind, feature_norm = self._gather(ind, feature_norm) 38 | feature_old = self.feature_bank[ind, ...] 39 | feature_new = (1 - self.momentum) * feature_old + \ 40 | self.momentum * feature_norm 41 | feature_new_norm = nn.functional.normalize(feature_new) 42 | self.feature_bank[ind, ...] = feature_new_norm 43 | 44 | def _gather(self, ind, feature): 45 | """Gather indices and features. 46 | 47 | Args: 48 | ind (Tensor): Indices for the batch of features. 49 | feature (Tensor): Batch of features. 50 | 51 | Returns: 52 | Tensor: Gathered indices. 53 | Tensor: Gathered features. 54 | """ 55 | ind_gathered = [ 56 | torch.ones_like(ind).cuda() for _ in range(self.num_replicas) 57 | ] 58 | feature_gathered = [ 59 | torch.ones_like(feature).cuda() for _ in range(self.num_replicas) 60 | ] 61 | dist.all_gather(ind_gathered, ind) 62 | dist.all_gather(feature_gathered, feature) 63 | ind_gathered = torch.cat(ind_gathered, dim=0) 64 | feature_gathered = torch.cat(feature_gathered, dim=0) 65 | return ind_gathered, feature_gathered 66 | -------------------------------------------------------------------------------- /openselfsup/models/registry.py: -------------------------------------------------------------------------------- 1 | from openselfsup.utils import Registry 2 | 3 | MODELS = Registry('model') 4 | BACKBONES = Registry('backbone') 5 | NECKS = Registry('neck') 6 | HEADS = Registry('head') 7 | MEMORIES = Registry('memory') 8 | LOSSES = Registry('loss') 9 | -------------------------------------------------------------------------------- /openselfsup/models/rotation_pred.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from openselfsup.utils import print_log 5 | 6 | from . import builder 7 | from .registry import MODELS 8 | 9 | 10 | @MODELS.register_module 11 | class RotationPred(nn.Module): 12 | """Rotation prediction. 13 | 14 | Implementation of "Unsupervised Representation Learning 15 | by Predicting Image Rotations (https://arxiv.org/abs/1803.07728)". 16 | 17 | Args: 18 | backbone (dict): Config dict for module of backbone ConvNet. 19 | head (dict): Config dict for module of loss functions. Default: None. 20 | pretrained (str, optional): Path to pre-trained weights. Default: None. 21 | """ 22 | 23 | def __init__(self, backbone, head=None, pretrained=None): 24 | super(RotationPred, self).__init__() 25 | self.backbone = builder.build_backbone(backbone) 26 | if head is not None: 27 | self.head = builder.build_head(head) 28 | self.init_weights(pretrained=pretrained) 29 | 30 | def init_weights(self, pretrained=None): 31 | """Initialize the weights of model. 32 | 33 | Args: 34 | pretrained (str, optional): Path to pre-trained weights. 35 | Default: None. 36 | """ 37 | if pretrained is not None: 38 | print_log('load model from: {}'.format(pretrained), logger='root') 39 | self.backbone.init_weights(pretrained=pretrained) 40 | self.head.init_weights(init_linear='kaiming') 41 | 42 | def forward_backbone(self, img): 43 | """Forward backbone. 44 | 45 | Args: 46 | img (Tensor): Input images of shape (N, C, H, W). 47 | Typically these should be mean centered and std scaled. 48 | 49 | Returns: 50 | tuple[Tensor]: backbone outputs. 51 | """ 52 | x = self.backbone(img) 53 | return x 54 | 55 | def forward_train(self, img, rot_label, **kwargs): 56 | """Forward computation during training. 57 | 58 | Args: 59 | img (Tensor): Input images of shape (N, C, H, W). 60 | Typically these should be mean centered and std scaled. 61 | rot_label (Tensor): Labels for the rotations. 62 | kwargs: Any keyword arguments to be used to forward. 63 | 64 | Returns: 65 | dict[str, Tensor]: A dictionary of loss components. 66 | """ 67 | x = self.forward_backbone(img) 68 | outs = self.head(x) 69 | loss_inputs = (outs, rot_label) 70 | losses = self.head.loss(*loss_inputs) 71 | return losses 72 | 73 | def forward_test(self, img, **kwargs): 74 | x = self.forward_backbone(img) # tuple 75 | outs = self.head(x) 76 | keys = ['head{}'.format(i) for i in range(len(outs))] 77 | out_tensors = [out.cpu() for out in outs] # NxC 78 | return dict(zip(keys, out_tensors)) 79 | 80 | def forward(self, img, rot_label=None, mode='train', **kwargs): 81 | if mode != "extract" and img.dim() == 5: # Nx4xCxHxW 82 | assert rot_label.dim() == 2 # Nx4 83 | img = img.view( 84 | img.size(0) * img.size(1), img.size(2), img.size(3), 85 | img.size(4)) # (4N)xCxHxW 86 | rot_label = torch.flatten(rot_label) # (4N) 87 | if mode == 'train': 88 | return self.forward_train(img, rot_label, **kwargs) 89 | elif mode == 'test': 90 | return self.forward_test(img, **kwargs) 91 | elif mode == 'extract': 92 | return self.forward_backbone(img) 93 | else: 94 | raise Exception("No such mode: {}".format(mode)) 95 | -------------------------------------------------------------------------------- /openselfsup/models/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .accuracy import Accuracy, accuracy 2 | from .conv_module import ConvModule, build_conv_layer 3 | from .conv_ws import ConvWS2d, conv_ws_2d 4 | from .gather_layer import GatherLayer 5 | from .multi_pooling import MultiPooling 6 | from .norm import build_norm_layer 7 | from .scale import Scale 8 | #from .weight_init import (bias_init_with_prob, kaiming_init, normal_init, 9 | # uniform_init, xavier_init) 10 | from .sobel import Sobel 11 | 12 | #__all__ = [ 13 | # 'conv_ws_2d', 'ConvWS2d', 'build_conv_layer', 'ConvModule', 14 | # 'build_norm_layer', 'xavier_init', 'normal_init', 'uniform_init', 15 | # 'kaiming_init', 'bias_init_with_prob', 'Scale', 'Sobel' 16 | #] 17 | -------------------------------------------------------------------------------- /openselfsup/models/utils/accuracy.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | def accuracy(pred, target, topk=1): 5 | assert isinstance(topk, (int, tuple)) 6 | if isinstance(topk, int): 7 | topk = (topk, ) 8 | return_single = True 9 | else: 10 | return_single = False 11 | 12 | maxk = max(topk) 13 | _, pred_label = pred.topk(maxk, dim=1) 14 | pred_label = pred_label.t() 15 | correct = pred_label.eq(target.view(1, -1).expand_as(pred_label)) 16 | 17 | res = [] 18 | for k in topk: 19 | correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) 20 | res.append(correct_k.mul_(100.0 / pred.size(0))) 21 | return res[0] if return_single else res 22 | 23 | 24 | class Accuracy(nn.Module): 25 | 26 | def __init__(self, topk=(1, )): 27 | super().__init__() 28 | self.topk = topk 29 | 30 | def forward(self, pred, target): 31 | return accuracy(pred, target, self.topk) 32 | -------------------------------------------------------------------------------- /openselfsup/models/utils/conv_ws.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | 4 | 5 | def conv_ws_2d(input, 6 | weight, 7 | bias=None, 8 | stride=1, 9 | padding=0, 10 | dilation=1, 11 | groups=1, 12 | eps=1e-5): 13 | c_in = weight.size(0) 14 | weight_flat = weight.view(c_in, -1) 15 | mean = weight_flat.mean(dim=1, keepdim=True).view(c_in, 1, 1, 1) 16 | std = weight_flat.std(dim=1, keepdim=True).view(c_in, 1, 1, 1) 17 | weight = (weight - mean) / (std + eps) 18 | return F.conv2d(input, weight, bias, stride, padding, dilation, groups) 19 | 20 | 21 | class ConvWS2d(nn.Conv2d): 22 | 23 | def __init__(self, 24 | in_channels, 25 | out_channels, 26 | kernel_size, 27 | stride=1, 28 | padding=0, 29 | dilation=1, 30 | groups=1, 31 | bias=True, 32 | eps=1e-5): 33 | super(ConvWS2d, self).__init__( 34 | in_channels, 35 | out_channels, 36 | kernel_size, 37 | stride=stride, 38 | padding=padding, 39 | dilation=dilation, 40 | groups=groups, 41 | bias=bias) 42 | self.eps = eps 43 | 44 | def forward(self, x): 45 | return conv_ws_2d(x, self.weight, self.bias, self.stride, self.padding, 46 | self.dilation, self.groups, self.eps) 47 | -------------------------------------------------------------------------------- /openselfsup/models/utils/gather_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.distributed as dist 3 | 4 | 5 | class GatherLayer(torch.autograd.Function): 6 | """Gather tensors from all process, supporting backward propagation. 7 | """ 8 | 9 | @staticmethod 10 | def forward(ctx, input): 11 | ctx.save_for_backward(input) 12 | output = [torch.zeros_like(input) \ 13 | for _ in range(dist.get_world_size())] 14 | dist.all_gather(output, input) 15 | return tuple(output) 16 | 17 | @staticmethod 18 | def backward(ctx, *grads): 19 | input, = ctx.saved_tensors 20 | grad_out = torch.zeros_like(input) 21 | grad_out[:] = grads[dist.get_rank()] 22 | return grad_out 23 | -------------------------------------------------------------------------------- /openselfsup/models/utils/multi_pooling.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class MultiPooling(nn.Module): 5 | """Pooling layers for features from multiple depth.""" 6 | 7 | POOL_PARAMS = { 8 | 'resnet50': [ 9 | dict(kernel_size=10, stride=10, padding=4), 10 | dict(kernel_size=16, stride=8, padding=0), 11 | dict(kernel_size=13, stride=5, padding=0), 12 | dict(kernel_size=8, stride=3, padding=0), 13 | dict(kernel_size=6, stride=1, padding=0) 14 | ] 15 | } 16 | POOL_SIZES = {'resnet50': [12, 6, 4, 3, 2]} 17 | POOL_DIMS = {'resnet50': [9216, 9216, 8192, 9216, 8192]} 18 | 19 | def __init__(self, 20 | pool_type='adaptive', 21 | in_indices=(0, ), 22 | backbone='resnet50'): 23 | super(MultiPooling, self).__init__() 24 | assert pool_type in ['adaptive', 'specified'] 25 | if pool_type == 'adaptive': 26 | self.pools = nn.ModuleList([ 27 | nn.AdaptiveAvgPool2d(self.POOL_SIZES[backbone][l]) 28 | for l in in_indices 29 | ]) 30 | else: 31 | self.pools = nn.ModuleList([ 32 | nn.AvgPool2d(**self.POOL_PARAMS[backbone][l]) 33 | for l in in_indices 34 | ]) 35 | 36 | def forward(self, x): 37 | assert isinstance(x, (list, tuple)) 38 | return [p(xx) for p, xx in zip(self.pools, x)] 39 | -------------------------------------------------------------------------------- /openselfsup/models/utils/norm.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | norm_cfg = { 4 | # format: layer_type: (abbreviation, module) 5 | 'BN': ('bn', nn.BatchNorm2d), 6 | 'SyncBN': ('bn', nn.SyncBatchNorm), 7 | 'GN': ('gn', nn.GroupNorm), 8 | # and potentially 'SN' 9 | } 10 | 11 | 12 | def build_norm_layer(cfg, num_features, postfix=''): 13 | """Build normalization layer. 14 | 15 | Args: 16 | cfg (dict): cfg should contain: 17 | type (str): identify norm layer type. 18 | layer args: args needed to instantiate a norm layer. 19 | requires_grad (bool): [optional] whether stop gradient updates 20 | num_features (int): number of channels from input. 21 | postfix (int, str): appended into norm abbreviation to 22 | create named layer. 23 | 24 | Returns: 25 | name (str): abbreviation + postfix 26 | layer (nn.Module): created norm layer 27 | """ 28 | assert isinstance(cfg, dict) and 'type' in cfg 29 | cfg_ = cfg.copy() 30 | 31 | layer_type = cfg_.pop('type') 32 | if layer_type not in norm_cfg: 33 | raise KeyError('Unrecognized norm type {}'.format(layer_type)) 34 | else: 35 | abbr, norm_layer = norm_cfg[layer_type] 36 | if norm_layer is None: 37 | raise NotImplementedError 38 | 39 | assert isinstance(postfix, (int, str)) 40 | name = abbr + str(postfix) 41 | 42 | requires_grad = cfg_.pop('requires_grad', True) 43 | cfg_.setdefault('eps', 1e-5) 44 | if layer_type != 'GN': 45 | layer = norm_layer(num_features, **cfg_) 46 | if layer_type == 'SyncBN': 47 | layer._specify_ddp_gpu_num(1) 48 | else: 49 | assert 'num_groups' in cfg_ 50 | layer = norm_layer(num_channels=num_features, **cfg_) 51 | 52 | for param in layer.parameters(): 53 | param.requires_grad = requires_grad 54 | 55 | return name, layer 56 | -------------------------------------------------------------------------------- /openselfsup/models/utils/scale.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class Scale(nn.Module): 6 | """A learnable scale parameter.""" 7 | 8 | def __init__(self, scale=1.0): 9 | super(Scale, self).__init__() 10 | self.scale = nn.Parameter(torch.tensor(scale, dtype=torch.float)) 11 | 12 | def forward(self, x): 13 | return x * self.scale 14 | -------------------------------------------------------------------------------- /openselfsup/models/utils/sobel.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class Sobel(nn.Module): 6 | """Sobel layer.""" 7 | 8 | def __init__(self): 9 | super(Sobel, self).__init__() 10 | grayscale = nn.Conv2d(3, 1, kernel_size=1, stride=1, padding=0) 11 | grayscale.weight.data.fill_(1.0 / 3.0) 12 | grayscale.bias.data.zero_() 13 | sobel_filter = nn.Conv2d(1, 2, kernel_size=3, stride=1, padding=1) 14 | sobel_filter.weight.data[0, 0].copy_( 15 | torch.FloatTensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]])) 16 | sobel_filter.weight.data[1, 0].copy_( 17 | torch.FloatTensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])) 18 | sobel_filter.bias.data.zero_() 19 | self.sobel = nn.Sequential(grayscale, sobel_filter) 20 | for p in self.sobel.parameters(): 21 | p.requires_grad = False 22 | 23 | def forward(self, x): 24 | return self.sobel(x) 25 | -------------------------------------------------------------------------------- /openselfsup/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .alias_multinomial import AliasMethod 2 | from .collect import nondist_forward_collect, dist_forward_collect 3 | from .collect_env import collect_env 4 | from .config_tools import traverse_replace 5 | from .flops_counter import get_model_complexity_info 6 | from .logger import get_root_logger, print_log 7 | from .registry import Registry, build_from_cfg 8 | from . import optimizers 9 | -------------------------------------------------------------------------------- /openselfsup/utils/alias_multinomial.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | class AliasMethod(object): 6 | """The alias method for sampling. 7 | 8 | From: https://hips.seas.harvard.edu/blog/2013/03/03/the-alias-method-efficient-sampling-with-many-discrete-outcomes/ 9 | 10 | Args: 11 | probs (Tensor): Sampling probabilities. 12 | """ 13 | 14 | def __init__(self, probs): 15 | 16 | if probs.sum() > 1: 17 | probs.div_(probs.sum()) 18 | K = len(probs) 19 | self.prob = torch.zeros(K) 20 | self.alias = torch.LongTensor([0] * K) 21 | 22 | # Sort the data into the outcomes with probabilities 23 | # that are larger and smaller than 1/K. 24 | smaller = [] 25 | larger = [] 26 | for kk, prob in enumerate(probs): 27 | self.prob[kk] = K * prob 28 | if self.prob[kk] < 1.0: 29 | smaller.append(kk) 30 | else: 31 | larger.append(kk) 32 | 33 | # Loop though and create little binary mixtures that 34 | # appropriately allocate the larger outcomes over the 35 | # overall uniform mixture. 36 | while len(smaller) > 0 and len(larger) > 0: 37 | small = smaller.pop() 38 | large = larger.pop() 39 | 40 | self.alias[small] = large 41 | self.prob[large] = (self.prob[large] - 1.0) + self.prob[small] 42 | 43 | if self.prob[large] < 1.0: 44 | smaller.append(large) 45 | else: 46 | larger.append(large) 47 | 48 | for last_one in smaller + larger: 49 | self.prob[last_one] = 1 50 | 51 | def cuda(self): 52 | self.prob = self.prob.cuda() 53 | self.alias = self.alias.cuda() 54 | 55 | def draw(self, N): 56 | """Draw N samples from multinomial. 57 | 58 | Args: 59 | N (int): Number of samples. 60 | 61 | Returns: 62 | Tensor: Samples. 63 | """ 64 | K = self.alias.size(0) 65 | 66 | kk = torch.zeros( 67 | N, dtype=torch.long, device=self.prob.device).random_(0, K) 68 | prob = self.prob.index_select(0, kk) 69 | alias = self.alias.index_select(0, kk) 70 | # b is whether a random number is greater than q 71 | b = torch.bernoulli(prob) 72 | oq = kk.mul(b.long()) 73 | oj = alias.mul((1 - b).long()) 74 | 75 | return oq + oj 76 | -------------------------------------------------------------------------------- /openselfsup/utils/collect.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import mmcv 4 | import torch 5 | 6 | from .gather import gather_tensors_batch 7 | 8 | 9 | def nondist_forward_collect(func, data_loader, length): 10 | """Forward and collect network outputs. 11 | 12 | This function performs forward propagation and collects outputs. 13 | It can be used to collect results, features, losses, etc. 14 | 15 | Args: 16 | func (function): The function to process data. The output must be 17 | a dictionary of CPU tensors. 18 | length (int): Expected length of output arrays. 19 | 20 | Returns: 21 | results_all (dict(np.ndarray)): The concatenated outputs. 22 | """ 23 | results = [] 24 | prog_bar = mmcv.ProgressBar(len(data_loader)) 25 | for i, data in enumerate(data_loader): 26 | with torch.no_grad(): 27 | result = func(**data) 28 | results.append(result) 29 | prog_bar.update() 30 | 31 | results_all = {} 32 | for k in results[0].keys(): 33 | results_all[k] = np.concatenate( 34 | [batch[k].numpy() for batch in results], axis=0) 35 | assert results_all[k].shape[0] == length 36 | return results_all 37 | 38 | 39 | def dist_forward_collect(func, data_loader, rank, length, ret_rank=-1): 40 | """Forward and collect network outputs in a distributed manner. 41 | 42 | This function performs forward propagation and collects outputs. 43 | It can be used to collect results, features, losses, etc. 44 | 45 | Args: 46 | func (function): The function to process data. The output must be 47 | a dictionary of CPU tensors. 48 | rank (int): This process id. 49 | length (int): Expected length of output arrays. 50 | ret_rank (int): The process that returns. 51 | Other processes will return None. 52 | 53 | Returns: 54 | results_all (dict(np.ndarray)): The concatenated outputs. 55 | """ 56 | results = [] 57 | if rank == 0: 58 | prog_bar = mmcv.ProgressBar(len(data_loader)) 59 | for idx, data in enumerate(data_loader): 60 | with torch.no_grad(): 61 | result = func(**data) # dict{key: tensor} 62 | results.append(result) 63 | 64 | if rank == 0: 65 | prog_bar.update() 66 | 67 | results_all = {} 68 | for k in results[0].keys(): 69 | results_cat = np.concatenate([batch[k].numpy() for batch in results], 70 | axis=0) 71 | if ret_rank == -1: 72 | results_gathered = gather_tensors_batch(results_cat, part_size=20) 73 | results_strip = np.concatenate(results_gathered, axis=0)[:length] 74 | else: 75 | results_gathered = gather_tensors_batch( 76 | results_cat, part_size=20, ret_rank=ret_rank) 77 | if rank == ret_rank: 78 | results_strip = np.concatenate( 79 | results_gathered, axis=0)[:length] 80 | else: 81 | results_strip = None 82 | results_all[k] = results_strip 83 | return results_all 84 | -------------------------------------------------------------------------------- /openselfsup/utils/collect_env.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import subprocess 3 | import sys 4 | from collections import defaultdict 5 | 6 | import cv2 7 | import mmcv 8 | import torch 9 | import torchvision 10 | 11 | import openselfsup 12 | 13 | 14 | def collect_env(): 15 | """Collect the information of the running environments.""" 16 | env_info = {} 17 | env_info['sys.platform'] = sys.platform 18 | env_info['Python'] = sys.version.replace('\n', '') 19 | 20 | cuda_available = torch.cuda.is_available() 21 | env_info['CUDA available'] = cuda_available 22 | 23 | if cuda_available: 24 | from torch.utils.cpp_extension import CUDA_HOME 25 | env_info['CUDA_HOME'] = CUDA_HOME 26 | 27 | if CUDA_HOME is not None and osp.isdir(CUDA_HOME): 28 | try: 29 | nvcc = osp.join(CUDA_HOME, 'bin/nvcc') 30 | nvcc = subprocess.check_output( 31 | '"{}" -V | tail -n1'.format(nvcc), shell=True) 32 | nvcc = nvcc.decode('utf-8').strip() 33 | except subprocess.SubprocessError: 34 | nvcc = 'Not Available' 35 | env_info['NVCC'] = nvcc 36 | 37 | devices = defaultdict(list) 38 | for k in range(torch.cuda.device_count()): 39 | devices[torch.cuda.get_device_name(k)].append(str(k)) 40 | for name, devids in devices.items(): 41 | env_info['GPU ' + ','.join(devids)] = name 42 | 43 | gcc = subprocess.check_output('gcc --version | head -n1', shell=True) 44 | gcc = gcc.decode('utf-8').strip() 45 | env_info['GCC'] = gcc 46 | 47 | env_info['PyTorch'] = torch.__version__ 48 | env_info['PyTorch compiling details'] = torch.__config__.show() 49 | 50 | env_info['TorchVision'] = torchvision.__version__ 51 | 52 | env_info['OpenCV'] = cv2.__version__ 53 | 54 | env_info['MMCV'] = mmcv.__version__ 55 | env_info['OpenSelfSup'] = openselfsup.__version__ 56 | #from openselfsup.ops import get_compiler_version, get_compiling_cuda_version 57 | #env_info['OpenSelfSup Compiler'] = get_compiler_version() 58 | #env_info['OpenSelfSup CUDA Compiler'] = get_compiling_cuda_version() 59 | return env_info 60 | 61 | 62 | if __name__ == "__main__": 63 | for name, val in collect_env().items(): 64 | print('{}: {}'.format(name, val)) 65 | -------------------------------------------------------------------------------- /openselfsup/utils/config_tools.py: -------------------------------------------------------------------------------- 1 | from mmcv import Config 2 | 3 | def traverse_replace(d, key, value): 4 | if isinstance(d, (dict, Config)): 5 | for k, v in d.items(): 6 | if k == key: 7 | d[k] = value 8 | else: 9 | traverse_replace(v, key, value) 10 | elif isinstance(d, (list, tuple, set)): 11 | for v in d: 12 | traverse_replace(v, key, value) 13 | -------------------------------------------------------------------------------- /openselfsup/utils/gather.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import torch 4 | import torch.distributed as dist 5 | 6 | 7 | def gather_tensors(input_array): 8 | world_size = dist.get_world_size() 9 | ## gather shapes first 10 | myshape = input_array.shape 11 | mycount = input_array.size 12 | shape_tensor = torch.Tensor(np.array(myshape)).cuda() 13 | all_shape = [ 14 | torch.Tensor(np.array(myshape)).cuda() for i in range(world_size) 15 | ] 16 | dist.all_gather(all_shape, shape_tensor) 17 | ## compute largest shapes 18 | all_shape = [x.cpu().numpy() for x in all_shape] 19 | all_count = [int(x.prod()) for x in all_shape] 20 | all_shape = [list(map(int, x)) for x in all_shape] 21 | max_count = max(all_count) 22 | ## padding tensors and gather them 23 | output_tensors = [ 24 | torch.Tensor(max_count).cuda() for i in range(world_size) 25 | ] 26 | padded_input_array = np.zeros(max_count) 27 | padded_input_array[:mycount] = input_array.reshape(-1) 28 | input_tensor = torch.Tensor(padded_input_array).cuda() 29 | dist.all_gather(output_tensors, input_tensor) 30 | ## unpadding gathered tensors 31 | padded_output = [x.cpu().numpy() for x in output_tensors] 32 | output = [ 33 | x[:all_count[i]].reshape(all_shape[i]) 34 | for i, x in enumerate(padded_output) 35 | ] 36 | return output 37 | 38 | 39 | def gather_tensors_batch(input_array, part_size=100, ret_rank=-1): 40 | # batch-wize gathering to avoid CUDA out of memory 41 | rank = dist.get_rank() 42 | all_features = [] 43 | part_num = input_array.shape[0] // part_size + 1 if input_array.shape[ 44 | 0] % part_size != 0 else input_array.shape[0] // part_size 45 | for i in range(part_num): 46 | part_feat = input_array[i * 47 | part_size:min((i + 1) * 48 | part_size, input_array.shape[0]), 49 | ...] 50 | assert part_feat.shape[ 51 | 0] > 0, "rank: {}, length of part features should > 0".format(rank) 52 | #print("rank: {}, gather part: {}/{}, length: {}".format(rank, i, part_num, len(part_feat))) 53 | gather_part_feat = gather_tensors(part_feat) 54 | all_features.append(gather_part_feat) 55 | if ret_rank == -1: 56 | all_features = [ 57 | np.concatenate([all_features[i][j] for i in range(part_num)], 58 | axis=0) for j in range(len(all_features[0])) 59 | ] 60 | return all_features 61 | else: 62 | if rank == ret_rank: 63 | all_features = [ 64 | np.concatenate([all_features[i][j] for i in range(part_num)], 65 | axis=0) for j in range(len(all_features[0])) 66 | ] 67 | return all_features 68 | else: 69 | return None 70 | -------------------------------------------------------------------------------- /openselfsup/utils/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from mmcv.runner import get_dist_info 4 | 5 | 6 | def get_root_logger(log_file=None, log_level=logging.INFO): 7 | """Get the root logger. 8 | 9 | The logger will be initialized if it has not been initialized. By default a 10 | StreamHandler will be added. If `log_file` is specified, a FileHandler will 11 | also be added. The name of the root logger is the top-level package name, 12 | e.g., "openselfsup". 13 | 14 | Args: 15 | log_file (str | None): The log filename. If specified, a FileHandler 16 | will be added to the root logger. 17 | log_level (int): The root logger level. Note that only the process of 18 | rank 0 is affected, while other processes will set the level to 19 | "Error" and be silent most of the time. 20 | 21 | Returns: 22 | logging.Logger: The root logger. 23 | """ 24 | logger = logging.getLogger(__name__.split('.')[0]) # i.e., openselfsup 25 | # if the logger has been initialized, just return it 26 | if logger.hasHandlers(): 27 | return logger 28 | 29 | format_str = '%(asctime)s - %(name)s - %(levelname)s - %(message)s' 30 | logging.basicConfig(format=format_str, level=log_level) 31 | rank, _ = get_dist_info() 32 | if rank != 0: 33 | logger.setLevel('ERROR') 34 | elif log_file is not None: 35 | file_handler = logging.FileHandler(log_file, 'w') 36 | file_handler.setFormatter(logging.Formatter(format_str)) 37 | file_handler.setLevel(log_level) 38 | logger.addHandler(file_handler) 39 | 40 | return logger 41 | 42 | 43 | def print_log(msg, logger=None, level=logging.INFO): 44 | """Print a log message. 45 | 46 | Args: 47 | msg (str): The message to be logged. 48 | logger (logging.Logger | str | None): The logger to be used. Some 49 | special loggers are: 50 | - "root": the root logger obtained with `get_root_logger()`. 51 | - "silent": no message will be printed. 52 | - None: The `print()` method will be used to print log messages. 53 | level (int): Logging level. Only available when `logger` is a Logger 54 | object or "root". 55 | """ 56 | if logger is None: 57 | print(msg) 58 | elif logger == 'root': 59 | _logger = get_root_logger() 60 | _logger.log(level, msg) 61 | elif isinstance(logger, logging.Logger): 62 | logger.log(level, msg) 63 | elif logger != 'silent': 64 | raise TypeError( 65 | 'logger should be either a logging.Logger object, "root", ' 66 | '"silent" or None, but got {}'.format(logger)) 67 | -------------------------------------------------------------------------------- /openselfsup/utils/misc.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import mmcv 4 | import numpy as np 5 | from six.moves import map, zip 6 | 7 | 8 | def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True): 9 | num_imgs = tensor.size(0) 10 | mean = np.array(mean, dtype=np.float32) 11 | std = np.array(std, dtype=np.float32) 12 | imgs = [] 13 | for img_id in range(num_imgs): 14 | img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0) 15 | img = mmcv.imdenormalize( 16 | img, mean, std, to_bgr=to_rgb).astype(np.uint8) 17 | imgs.append(np.ascontiguousarray(img)) 18 | return imgs 19 | 20 | 21 | def multi_apply(func, *args, **kwargs): 22 | pfunc = partial(func, **kwargs) if kwargs else func 23 | map_results = map(pfunc, *args) 24 | return tuple(map(list, zip(*map_results))) 25 | 26 | 27 | def unmap(data, count, inds, fill=0): 28 | """Unmap a subset of item (data) back to the original set of items (of 29 | size count).""" 30 | if data.dim() == 1: 31 | ret = data.new_full((count, ), fill) 32 | ret[inds] = data 33 | else: 34 | new_size = (count, ) + data.size()[1:] 35 | ret = data.new_full(new_size, fill) 36 | ret[inds, :] = data 37 | return ret 38 | -------------------------------------------------------------------------------- /openselfsup/utils/profiling.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import sys 3 | import time 4 | 5 | import torch 6 | 7 | if sys.version_info >= (3, 7): 8 | 9 | @contextlib.contextmanager 10 | def profile_time(trace_name, 11 | name, 12 | enabled=True, 13 | stream=None, 14 | end_stream=None): 15 | """Print time spent by CPU and GPU. 16 | 17 | Useful as a temporary context manager to find sweet spots of 18 | code suitable for async implementation. 19 | """ 20 | if (not enabled) or not torch.cuda.is_available(): 21 | yield 22 | return 23 | stream = stream if stream else torch.cuda.current_stream() 24 | end_stream = end_stream if end_stream else stream 25 | start = torch.cuda.Event(enable_timing=True) 26 | end = torch.cuda.Event(enable_timing=True) 27 | stream.record_event(start) 28 | try: 29 | cpu_start = time.monotonic() 30 | yield 31 | finally: 32 | cpu_end = time.monotonic() 33 | end_stream.record_event(end) 34 | end.synchronize() 35 | cpu_time = (cpu_end - cpu_start) * 1000 36 | gpu_time = start.elapsed_time(end) 37 | msg = "{} {} cpu_time {:.2f} ms ".format(trace_name, name, 38 | cpu_time) 39 | msg += "gpu_time {:.2f} ms stream {}".format(gpu_time, stream) 40 | print(msg, end_stream) 41 | -------------------------------------------------------------------------------- /openselfsup/utils/registry.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | from functools import partial 3 | 4 | import mmcv 5 | 6 | 7 | class Registry(object): 8 | 9 | def __init__(self, name): 10 | self._name = name 11 | self._module_dict = dict() 12 | 13 | def __repr__(self): 14 | format_str = self.__class__.__name__ + '(name={}, items={})'.format( 15 | self._name, list(self._module_dict.keys())) 16 | return format_str 17 | 18 | @property 19 | def name(self): 20 | return self._name 21 | 22 | @property 23 | def module_dict(self): 24 | return self._module_dict 25 | 26 | def get(self, key): 27 | return self._module_dict.get(key, None) 28 | 29 | def _register_module(self, module_class, force=False): 30 | """Register a module. 31 | 32 | Args: 33 | module (:obj:`nn.Module`): Module to be registered. 34 | """ 35 | if not inspect.isclass(module_class): 36 | raise TypeError('module must be a class, but got {}'.format( 37 | type(module_class))) 38 | module_name = module_class.__name__ 39 | if not force and module_name in self._module_dict: 40 | raise KeyError('{} is already registered in {}'.format( 41 | module_name, self.name)) 42 | self._module_dict[module_name] = module_class 43 | 44 | def register_module(self, cls=None, force=False): 45 | if cls is None: 46 | return partial(self.register_module, force=force) 47 | self._register_module(cls, force=force) 48 | return cls 49 | 50 | 51 | def build_from_cfg(cfg, registry, default_args=None): 52 | """Build a module from config dict. 53 | 54 | Args: 55 | cfg (dict): Config dict. It should at least contain the key "type". 56 | registry (:obj:`Registry`): The registry to search the type from. 57 | default_args (dict, optional): Default initialization arguments. 58 | 59 | Returns: 60 | obj: The constructed object. 61 | """ 62 | assert isinstance(cfg, dict) and 'type' in cfg 63 | assert isinstance(default_args, dict) or default_args is None 64 | args = cfg.copy() 65 | obj_type = args.pop('type') 66 | if mmcv.is_str(obj_type): 67 | obj_cls = registry.get(obj_type) 68 | if obj_cls is None: 69 | raise KeyError('{} is not in the {} registry'.format( 70 | obj_type, registry.name)) 71 | elif inspect.isclass(obj_type): 72 | obj_cls = obj_type 73 | else: 74 | raise TypeError('type must be a str or valid type, but got {}'.format( 75 | type(obj_type))) 76 | if default_args is not None: 77 | for name, value in default_args.items(): 78 | args.setdefault(name, value) 79 | return obj_cls(**args) 80 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | -r requirements/runtime.txt 2 | -r requirements/tests.txt 3 | -------------------------------------------------------------------------------- /requirements/runtime.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | mmcv==1.0.3 3 | numpy 4 | # need older pillow until torchvision is fixed 5 | Pillow<=6.2.2 6 | packaging 7 | six 8 | terminaltables 9 | sklearn 10 | faiss-gpu==1.6.1 11 | tensorboard 12 | future 13 | tqdm 14 | -------------------------------------------------------------------------------- /requirements/tests.txt: -------------------------------------------------------------------------------- 1 | asynctest 2 | codecov 3 | flake8 4 | isort 5 | pytest 6 | pytest-cov 7 | pytest-runner 8 | xdoctest >= 0.10.0 9 | yapf 10 | # Note: used for kwarray.group_items, this may be ported to mmcv in the future. 11 | kwarray 12 | -------------------------------------------------------------------------------- /tools/count_parameters.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from mmcv import Config 3 | 4 | from openselfsup.models import build_model 5 | 6 | 7 | def parse_args(): 8 | parser = argparse.ArgumentParser(description='Train a model') 9 | parser.add_argument('config', help='train config file path') 10 | args = parser.parse_args() 11 | return args 12 | 13 | 14 | def main(): 15 | args = parse_args() 16 | 17 | cfg = Config.fromfile(args.config) 18 | 19 | model = build_model(cfg.model) 20 | 21 | num_params = sum(p.numel() for p in model.parameters()) / 1000000. 22 | num_grad_params = sum(p.numel() for p in model.parameters() \ 23 | if p.requires_grad) / 1000000. 24 | num_backbone_params = sum( 25 | p.numel() for p in model.backbone.parameters()) / 1000000. 26 | num_backbone_grad_params = sum(p.numel() for p in model.backbone.parameters() \ 27 | if p.requires_grad) / 1000000. 28 | print( 29 | "Number of backbone parameters: {:.5g} M".format(num_backbone_params)) 30 | print("Number of backbone parameters requiring grad: {:.5g} M".format( 31 | num_backbone_grad_params)) 32 | print("Number of total parameters: {:.5g} M".format(num_params)) 33 | print("Number of total parameters requiring grad: {:.5g} M".format( 34 | num_grad_params)) 35 | 36 | 37 | if __name__ == '__main__': 38 | main() 39 | -------------------------------------------------------------------------------- /tools/dist_extract.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -x 3 | 4 | PYTHON=${PYTHON:-"python"} 5 | CFG=$1 6 | GPUS=$2 7 | WORK_DIR=$3 8 | PY_ARGS=${@:4} # "--checkpoint $CHECKPOINT --pretrained $PRETRAINED" 9 | PORT=${PORT:-29500} 10 | 11 | $PYTHON -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 12 | tools/extract.py $CFG --layer-ind "0,1,2,3,4" --work_dir $WORK_DIR \ 13 | --launcher pytorch ${PY_ARGS} 14 | -------------------------------------------------------------------------------- /tools/dist_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -x 4 | 5 | CFG=$1 6 | GPUS=$2 7 | CHECKPOINT=$3 8 | PORT=${PORT:-29500} 9 | 10 | WORK_DIR="$(dirname $CHECKPOINT)/" 11 | 12 | # test 13 | python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 14 | tools/test.py \ 15 | $CFG \ 16 | $CHECKPOINT \ 17 | --work_dir $WORK_DIR --launcher="pytorch" 18 | -------------------------------------------------------------------------------- /tools/dist_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | PYTHON=${PYTHON:-"python"} 3 | 4 | CFG=$1 5 | GPUS=$2 6 | PY_ARGS=${@:3} 7 | PORT=${PORT:-29500} 8 | 9 | WORK_DIR=$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/ 10 | 11 | $PYTHON -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 12 | tools/train.py $CFG --work_dir $WORK_DIR --seed 0 --launcher pytorch ${PY_ARGS} 13 | -------------------------------------------------------------------------------- /tools/extract_backbone_weights.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import argparse 3 | 4 | 5 | def parse_args(): 6 | parser = argparse.ArgumentParser( 7 | description='This script extracts backbone weights from a checkpoint') 8 | parser.add_argument('checkpoint', help='checkpoint file') 9 | parser.add_argument( 10 | 'output', type=str, help='destination file name') 11 | args = parser.parse_args() 12 | return args 13 | 14 | 15 | def main(): 16 | args = parse_args() 17 | assert args.output.endswith(".pth") 18 | ck = torch.load(args.checkpoint, map_location=torch.device('cpu')) 19 | output_dict = dict(state_dict=dict(), author="OpenSelfSup") 20 | has_backbone = False 21 | for key, value in ck['state_dict'].items(): 22 | if key.startswith('backbone'): 23 | output_dict['state_dict'][key[9:]] = value 24 | has_backbone = True 25 | if not has_backbone: 26 | raise Exception("Cannot find a backbone module in the checkpoint.") 27 | torch.save(output_dict, args.output) 28 | 29 | 30 | if __name__ == '__main__': 31 | main() 32 | -------------------------------------------------------------------------------- /tools/kill.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | kill $(ps aux | grep "train.py" | grep -v grep | awk '{print $2}') 3 | -------------------------------------------------------------------------------- /tools/prepare_data/convert_subset.py: -------------------------------------------------------------------------------- 1 | """ 2 | SimCLR provides list files for semi-supervised benchmarks: 3 | https://github.com/google-research/simclr/tree/master/imagenet_subsets/ 4 | This script convert the list files into the required format in OpenSelfSup. 5 | """ 6 | import argparse 7 | 8 | parser = argparse.ArgumentParser( 9 | description='Convert ImageNet subset lists provided by simclr.') 10 | parser.add_argument('input', help='Input list file.') 11 | parser.add_argument('output', help='Output list file.') 12 | args = parser.parse_args() 13 | 14 | # create dict 15 | with open("data/imagenet/meta/train_labeled.txt", 'r') as f: 16 | lines = f.readlines() 17 | keys = [l.split('/')[0] for l in lines] 18 | labels = [l.strip().split()[1] for l in lines] 19 | mapping = {} 20 | for k,l in zip(keys, labels): 21 | if k not in mapping: 22 | mapping[k] = l 23 | else: 24 | assert mapping[k] == l 25 | 26 | # convert 27 | with open(args.input, 'r') as f: 28 | lines = f.readlines() 29 | fns = [l.strip() for l in lines] 30 | sample_keys = [l.split('_')[0] for l in lines] 31 | sample_labels = [mapping[k] for k in sample_keys] 32 | output_lines = ["{}/{} {}\n".format(k, fn, l) for \ 33 | k,fn,l in zip(sample_keys, fns, sample_labels)] 34 | with open(args.output, 'w') as f: 35 | f.writelines(output_lines) 36 | -------------------------------------------------------------------------------- /tools/prepare_data/prepare_voc07_cls.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | DATA="$1" 4 | if [ "$DATA" == "" ]; then 5 | echo "Usage: bash tools/prepare_data/prepare_voc07_cls.sh YOUR_DATA_ROOT" 6 | exit 7 | fi 8 | 9 | VOC="$DATA/VOCdevkit/VOC2007/" 10 | 11 | wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar -P $DATA 12 | wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar -P $DATA 13 | tar -xf $DATA/VOCtrainval_06-Nov-2007.tar -C $DATA 14 | tar -xf $DATA/VOCtest_06-Nov-2007.tar -C $DATA 15 | 16 | mkdir -p $VOC/SVMLabels/low_shot/labels/ 17 | 18 | python $(dirname "$0")/create_voc_data_files.py \ 19 | --data_source_dir $VOC \ 20 | --output_dir $VOC/SVMLabels/ \ 21 | --generate_json 1 22 | 23 | python $(dirname "$0")/create_voc_low_shot_challenge_samples.py \ 24 | --targets_data_file $VOC/SVMLabels/train_targets.json \ 25 | --output_path $VOC/SVMLabels/low_shot/labels/ \ 26 | --k_values "1,2,4,8,16,32,64,96" \ 27 | --num_samples 5 28 | 29 | mkdir $VOC/Lists 30 | 31 | awk 'NF{print $0 ".jpg"}' $VOC/ImageSets/Main/trainval.txt $VOC/ImageSets/Main/test.txt > $VOC/Lists/trainvaltest.txt 32 | 33 | mkdir data/ 34 | ln -s $DATA/VOCdevkit data/ 35 | -------------------------------------------------------------------------------- /tools/publish_model.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import subprocess 3 | 4 | 5 | def parse_args(): 6 | parser = argparse.ArgumentParser( 7 | description='Process a checkpoint to be published') 8 | parser.add_argument('in_file', help='input checkpoint filename') 9 | args = parser.parse_args() 10 | return args 11 | 12 | 13 | def process_checkpoint(in_file): 14 | tmp_file = in_file + ".tmp" 15 | subprocess.Popen(['cp', in_file, tmp_file]) 16 | sha = subprocess.check_output(['sha256sum', tmp_file]).decode() 17 | out_file = in_file 18 | if out_file.endswith('.pth'): 19 | out_file = out_file[:-4] 20 | final_file = out_file + f'-{sha[:8]}.pth' 21 | assert final_file != in_file, \ 22 | "The output filename is the same as the input file." 23 | print("Output file: {}".format(final_file)) 24 | subprocess.Popen(['mv', tmp_file, final_file]) 25 | 26 | 27 | def main(): 28 | args = parse_args() 29 | process_checkpoint(args.in_file) 30 | 31 | 32 | if __name__ == '__main__': 33 | main() 34 | -------------------------------------------------------------------------------- /tools/single_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | PYTHON=${PYTHON:-"python"} 3 | 4 | CFG=$1 5 | PY_ARGS=${@:2} 6 | 7 | WORK_DIR=$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/ 8 | 9 | $PYTHON -u tools/train.py $1 --work_dir $WORK_DIR ${PY_ARGS} 10 | -------------------------------------------------------------------------------- /tools/srun_extract.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -x 3 | 4 | PARTITION=$1 5 | CFG=$2 6 | GPUS=$3 7 | WORK_DIR=$4 8 | PY_ARGS=${@:5} # "--checkpoint $CHECKPOINT --pretrained $PRETRAINED" 9 | JOB_NAME="openselfsup" 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | SRUN_ARGS=${SRUN_ARGS:-""} 13 | 14 | srun -p ${PARTITION} \ 15 | --job-name=${JOB_NAME} \ 16 | --gres=gpu:${GPUS_PER_NODE} \ 17 | --ntasks=${GPUS} \ 18 | --ntasks-per-node=${GPUS_PER_NODE} \ 19 | --cpus-per-task=${CPUS_PER_TASK} \ 20 | --kill-on-bad-exit=1 \ 21 | ${SRUN_ARGS} \ 22 | python -u tools/extract.py $CFG \ 23 | --layer-ind "0,1,2,3,4" --work_dir $WORK_DIR \ 24 | --launcher="slurm" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /tools/srun_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | CFG=$2 7 | GPUS=$3 8 | CHECKPOINT=$4 9 | PY_ARGS=${@:5} # --port 10 | JOB_NAME="openselfsup" 11 | GPUS_PER_NODE=${GPUS_PER_NODE:-1} 12 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 13 | SRUN_ARGS=${SRUN_ARGS:-""} 14 | 15 | WORK_DIR="$(dirname $CHECKPOINT)/" 16 | 17 | # test 18 | GLOG_vmodule=MemcachedClient=-1 \ 19 | srun -p ${PARTITION} \ 20 | --job-name=${JOB_NAME} \ 21 | --gres=gpu:${GPUS_PER_NODE} \ 22 | --ntasks=${GPUS} \ 23 | --ntasks-per-node=${GPUS_PER_NODE} \ 24 | --cpus-per-task=${CPUS_PER_TASK} \ 25 | --kill-on-bad-exit=1 \ 26 | ${SRUN_ARGS} \ 27 | python -u tools/test.py \ 28 | $CFG \ 29 | $CHECKPOINT \ 30 | --work_dir $WORK_DIR --launcher="slurm" $PY_ARGS 31 | -------------------------------------------------------------------------------- /tools/srun_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | CFG=$2 7 | GPUS=$3 8 | PY_ARGS=${@:4} 9 | JOB_NAME="openselfsup" 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | SRUN_ARGS=${SRUN_ARGS:-""} 13 | 14 | WORK_DIR=$(echo ${CFG%.*} | sed -e "s/configs/work_dirs/g")/ 15 | 16 | GLOG_vmodule=MemcachedClient=-1 \ 17 | srun -p ${PARTITION} \ 18 | --job-name=${JOB_NAME} \ 19 | --gres=gpu:${GPUS_PER_NODE} \ 20 | --ntasks=${GPUS} \ 21 | --ntasks-per-node=${GPUS_PER_NODE} \ 22 | --cpus-per-task=${CPUS_PER_TASK} \ 23 | --kill-on-bad-exit=1 \ 24 | ${SRUN_ARGS} \ 25 | python -u tools/train.py ${CFG} \ 26 | --work_dir ${WORK_DIR} --seed 0 --launcher="slurm" ${PY_ARGS} 27 | -------------------------------------------------------------------------------- /tools/upgrade_models.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import argparse 3 | 4 | 5 | def parse_args(): 6 | parser = argparse.ArgumentParser() 7 | parser.add_argument('checkpoint', help='checkpoint file') 8 | parser.add_argument( 9 | '--save-path', type=str, required=True, help='destination file name') 10 | args = parser.parse_args() 11 | return args 12 | 13 | 14 | def main(): 15 | args = parse_args() 16 | ck = torch.load(args.checkpoint, map_location=torch.device('cpu')) 17 | output_dict = dict(state_dict=dict(), author='OpenSelfSup') 18 | for key, value in ck.items(): 19 | if key.startswith('head'): 20 | continue 21 | else: 22 | output_dict['state_dict'][key] = value 23 | torch.save(output_dict, args.save_path) 24 | 25 | 26 | if __name__ == '__main__': 27 | main() 28 | --------------------------------------------------------------------------------