├── src ├── utils │ ├── __init__.py │ ├── data_util.py │ ├── main_util.py │ ├── misc_util.py │ └── coco_util.py ├── distillation │ ├── __init__.py │ ├── loss.py │ └── tool.py ├── models │ ├── mimic │ │ ├── __init__.py │ │ ├── base.py │ │ └── resnet_layer.py │ ├── org │ │ └── __init__.py │ ├── custom │ │ ├── __init__.py │ │ └── resnet.py │ ├── ext │ │ ├── __init__.py │ │ ├── classifier.py │ │ └── backbone.py │ └── __init__.py ├── structure │ ├── __init__.py │ ├── transformer.py │ └── sampler.py ├── visualizer.py ├── coco_runner.py ├── mimic_runner.py └── ext_runner.py ├── .gitignore ├── img └── ghnd_and_nf.png ├── .gitmodules ├── Pipfile ├── LICENSE ├── config ├── org │ ├── faster_rcnn-backbone_resnet50.yaml │ ├── mask_rcnn-backbone_resnet50.yaml │ └── keypoint_rcnn-backbone_resnet50.yaml ├── ext │ └── keypoint_rcnn-backbone_ext_resnet50-b3ch.yaml ├── hnd │ ├── mask_rcnn-backbone_resnet50-b12ch.yaml │ ├── mask_rcnn-backbone_resnet50-b15ch.yaml │ ├── mask_rcnn-backbone_resnet50-b3ch.yaml │ ├── mask_rcnn-backbone_resnet50-b6ch.yaml │ ├── mask_rcnn-backbone_resnet50-b9ch.yaml │ ├── faster_rcnn-backbone_resnet50-b12ch.yaml │ ├── faster_rcnn-backbone_resnet50-b15ch.yaml │ ├── faster_rcnn-backbone_resnet50-b3ch.yaml │ ├── faster_rcnn-backbone_resnet50-b6ch.yaml │ ├── faster_rcnn-backbone_resnet50-b9ch.yaml │ ├── keypoint_rcnn-backbone_resnet50-b3ch.yaml │ ├── keypoint_rcnn-backbone_resnet50-b6ch.yaml │ ├── keypoint_rcnn-backbone_resnet50-b9ch.yaml │ ├── keypoint_rcnn-backbone_resnet50-b12ch.yaml │ └── keypoint_rcnn-backbone_resnet50-b15ch.yaml └── ghnd │ ├── mask_rcnn-backbone_resnet50-b12ch.yaml │ ├── mask_rcnn-backbone_resnet50-b3ch.yaml │ ├── mask_rcnn-backbone_resnet50-b6ch.yaml │ ├── mask_rcnn-backbone_resnet50-b9ch.yaml │ ├── faster_rcnn-backbone_resnet50-b12ch.yaml │ ├── faster_rcnn-backbone_resnet50-b3ch.yaml │ ├── faster_rcnn-backbone_resnet50-b6ch.yaml │ ├── faster_rcnn-backbone_resnet50-b9ch.yaml │ ├── keypoint_rcnn-backbone_resnet50-b12ch.yaml │ ├── keypoint_rcnn-backbone_resnet50-b3ch.yaml │ ├── keypoint_rcnn-backbone_resnet50-b6ch.yaml │ └── keypoint_rcnn-backbone_resnet50-b9ch.yaml └── README.md /src/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/distillation/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/mimic/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/org/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/structure/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/custom/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | .ipynb_checkpoints/ 3 | __pycache__/ 4 | resource/ 5 | .editorconfig 6 | -------------------------------------------------------------------------------- /img/ghnd_and_nf.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yoshitomo-matsubara/hnd-ghnd-object-detectors/HEAD/img/ghnd_and_nf.png -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "src/myutils"] 2 | path = src/myutils 3 | url = https://github.com/yoshitomo-matsubara/myutils.git 4 | -------------------------------------------------------------------------------- /Pipfile: -------------------------------------------------------------------------------- 1 | [[source]] 2 | url = "https://pypi.org/simple" 3 | verify_ssl = true 4 | name = "pypi" 5 | 6 | [packages] 7 | torch = "==1.3.1" 8 | torchvision = "==0.4.2" 9 | pillow = "<7" 10 | numpy = "*" 11 | matplotlib = "*" 12 | pyyaml = "*" 13 | pycocotools = "==2.0.1" 14 | cython = "*" 15 | scikit-learn = "*" 16 | pandas = "*" 17 | opencv-python = "*" 18 | 19 | [dev-packages] 20 | 21 | [requires] 22 | python_version = "3.6" 23 | -------------------------------------------------------------------------------- /src/models/ext/__init__.py: -------------------------------------------------------------------------------- 1 | from models.ext.backbone import ExtBackboneWithFPN 2 | 3 | 4 | def get_ext_fpn_backbone(base_backbone, ext_config, freeze_layers): 5 | if freeze_layers: 6 | for name, parameter in base_backbone.named_parameters(): 7 | if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: 8 | parameter.requires_grad_(False) 9 | 10 | return_layers = {'layer1': 0, 'layer2': 1, 'layer3': 2, 'layer4': 3} 11 | in_channels_stage2 = base_backbone.inplanes // 8 12 | in_channels_list = [ 13 | in_channels_stage2, 14 | in_channels_stage2 * 2, 15 | in_channels_stage2 * 4, 16 | in_channels_stage2 * 8, 17 | ] 18 | out_channels = 256 19 | return ExtBackboneWithFPN(base_backbone, return_layers, in_channels_list, out_channels, ext_config) 20 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Yoshitomo Matsubara 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /src/models/ext/classifier.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torchvision.models.resnet import ResNet 3 | 4 | from models.custom.resnet import CustomResNet 5 | 6 | 7 | class BaseExtClassifier(nn.Module): 8 | def __init__(self, ext_idx): 9 | super().__init__() 10 | self.ext_idx = ext_idx 11 | 12 | def forward(self, *args): 13 | raise NotImplementedError('forward function is not implemented') 14 | 15 | 16 | class Ext4ResNet(BaseExtClassifier): 17 | def __init__(self, input_channel): 18 | super().__init__(ext_idx=0) 19 | self.extractor = nn.Sequential( 20 | nn.AdaptiveAvgPool2d((64, 64)), 21 | nn.Conv2d(input_channel, 64, kernel_size=4, stride=2), 22 | nn.BatchNorm2d(64), 23 | nn.ReLU(inplace=True), 24 | nn.Conv2d(64, 32, kernel_size=3, stride=2), 25 | nn.BatchNorm2d(32), 26 | nn.ReLU(inplace=True), 27 | nn.Conv2d(32, 16, kernel_size=2, stride=1), 28 | nn.BatchNorm2d(16), 29 | nn.ReLU(inplace=True), 30 | nn.AdaptiveAvgPool2d((8, 8)) 31 | ) 32 | self.linear = nn.Linear(16 * 8 * 8, 2) 33 | 34 | def forward(self, x): 35 | z = self.extractor(x) 36 | z = self.linear(z.flatten(1)) 37 | return z if self.training else z.softmax(dim=1) 38 | 39 | 40 | def get_ext_classifier(backbone): 41 | if isinstance(backbone, (ResNet, CustomResNet)): 42 | return Ext4ResNet(backbone) 43 | raise ValueError('type of backbone `{}` is not expected'.format(type(backbone))) 44 | -------------------------------------------------------------------------------- /config/org/faster_rcnn-backbone_resnet50.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | model: 24 | name: &model_name 'faster_rcnn' 25 | backbone: 26 | name: &backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &experiment !join [*dataset_name, '-', *model_name, '-backbone_', *backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *experiment, '.pt'] 35 | 36 | train: 37 | num_epochs: 26 38 | batch_size: 2 39 | log_freq: 1000 40 | optimizer: 41 | type: 'SGD' 42 | params: 43 | lr: 0.0075 # 0.02 / 8 * num_gpus 44 | momentum: 0.9 45 | weight_decay: 0.0001 46 | scheduler: 47 | type: 'MultiStepLR' 48 | params: 49 | milestones: [16, 22] 50 | gamma: 0.1 51 | 52 | test: 53 | batch_size: 1 54 | -------------------------------------------------------------------------------- /config/org/mask_rcnn-backbone_resnet50.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | model: 24 | name: &model_name 'mask_rcnn' 25 | backbone: 26 | name: &backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &experiment !join [*dataset_name, '-', *model_name, '-backbone_', *backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *experiment, '.pt'] 35 | 36 | train: 37 | num_epochs: 26 38 | batch_size: 2 39 | log_freq: 1000 40 | optimizer: 41 | type: 'SGD' 42 | params: 43 | lr: 0.0075 # 0.02 / 8 * num_gpus 44 | momentum: 0.9 45 | weight_decay: 0.0001 46 | scheduler: 47 | type: 'MultiStepLR' 48 | params: 49 | milestones: [16, 22] 50 | gamma: 0.1 51 | 52 | test: 53 | batch_size: 1 54 | -------------------------------------------------------------------------------- /config/org/keypoint_rcnn-backbone_resnet50.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | model: 24 | name: &model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &experiment !join [*dataset_name, '-', *model_name, '-backbone_', *backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *experiment, '.pt'] 36 | 37 | train: 38 | num_epochs: 46 39 | batch_size: 2 40 | log_freq: 1000 41 | optimizer: 42 | type: 'SGD' 43 | params: 44 | lr: 0.0075 # 0.02 / 8 * num_gpus 45 | momentum: 0.9 46 | weight_decay: 0.0001 47 | scheduler: 48 | type: 'MultiStepLR' 49 | params: 50 | milestones: [36, 43] 51 | gamma: 0.1 52 | 53 | test: 54 | batch_size: 1 55 | -------------------------------------------------------------------------------- /src/distillation/loss.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from myutils.pytorch import func_util 4 | 5 | 6 | class CustomLoss(nn.Module): 7 | def __init__(self, criterion_config): 8 | super().__init__() 9 | self.org_loss_factor = criterion_config['params']['org_loss_factor'] 10 | term_dict = dict() 11 | for loss_name, loss_config in criterion_config['terms'].items(): 12 | sub_criterion_config = loss_config['criterion'] 13 | sub_criterion = func_util.get_loss(sub_criterion_config['type'], sub_criterion_config['params']) 14 | term_dict[loss_name] = (loss_config['ts_modules'], sub_criterion, loss_config['factor']) 15 | self.term_dict = term_dict 16 | 17 | def forward(self, *args, **kwargs): 18 | raise NotImplementedError('forward function is not implemented') 19 | 20 | 21 | class GeneralizedCustomLoss(CustomLoss): 22 | def __init__(self, criterion_config): 23 | super().__init__(criterion_config) 24 | 25 | def forward(self, output_dict, org_loss_dict): 26 | loss_dict = dict() 27 | for loss_name, ((teacher_path, teacher_output), (student_path, student_output)) in output_dict.items(): 28 | _, criterion, factor = self.term_dict[loss_name] 29 | loss_dict[loss_name] = criterion(teacher_output, student_output) * factor 30 | 31 | sub_total_loss = sum(loss for loss in loss_dict.values()) 32 | if self.org_loss_factor == 0: 33 | return sub_total_loss 34 | return sub_total_loss + self.org_loss_factor * sum(loss for loss in org_loss_dict.values()) 35 | 36 | 37 | LOSS_DICT = { 38 | 'general': GeneralizedCustomLoss 39 | } 40 | 41 | 42 | def get_loss(criterion_config): 43 | criterion_type = criterion_config['type'] 44 | if criterion_type in LOSS_DICT: 45 | return LOSS_DICT[criterion_type](criterion_config) 46 | raise ValueError('criterion_type `{}` is not expected'.format(criterion_type)) 47 | -------------------------------------------------------------------------------- /src/models/mimic/base.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from structure.transformer import DataLogger 4 | 5 | 6 | class ExtEncoder(nn.Module): 7 | def __init__(self, encoder, ext_classifier=None, ext_config=None): 8 | super().__init__() 9 | self.encoder = encoder 10 | self.ext_classifier = ext_classifier 11 | self.threshold = ext_config['threshold'] if ext_config is not None else None 12 | 13 | def forward_with_ext(self, x): 14 | ext_z = self.ext_classifier(x) 15 | if not self.training and ext_z.shape[0] == 1 and ext_z[0][1] < self.threshold: 16 | return None, ext_z 17 | 18 | z = self.encoder(x) 19 | return z, ext_z 20 | 21 | def forward(self, x): 22 | return self.encoder(x) if self.ext_classifier is None else self.forward_with_ext(x) 23 | 24 | def get_ext_classifier(self): 25 | return self.ext_classifier 26 | 27 | 28 | class BottleneckBase4Ext(nn.Module): 29 | def __init__(self, encoder, decoder, bottleneck_transformer=None): 30 | super().__init__() 31 | self.encoder = encoder 32 | self.decoder = decoder 33 | self.bottleneck_transformer = bottleneck_transformer 34 | self.data_logging = isinstance(self.bottleneck_transformer, DataLogger) 35 | self.uses_ext_encoder = isinstance(encoder, ExtEncoder) and encoder.ext_classifier is not None 36 | self.use_bottleneck_transformer = False 37 | 38 | def forward_ext(self, z): 39 | z, ext_z = z 40 | if z is None: 41 | if self.data_logging: 42 | self.bottleneck_transformer(None, target=None) 43 | return z, ext_z 44 | elif not self.training and self.bottleneck_transformer is not None and self.use_bottleneck_transformer: 45 | device = z.device 46 | z, _ = self.bottleneck_transformer(z, target=None) 47 | z = z.to(device) 48 | return self.decoder(z), ext_z 49 | 50 | def forward(self, x): 51 | z = self.encoder(x) 52 | if self.uses_ext_encoder: 53 | return self.forward_ext(z) 54 | elif not self.training and self.bottleneck_transformer is not None and self.use_bottleneck_transformer: 55 | device = z.device 56 | z, _ = self.bottleneck_transformer(z, target=None) 57 | z = z.to(device) 58 | return self.decoder(z) 59 | 60 | def get_ext_classifier(self): 61 | raise NotImplementedError('get_ext_classifier function is not implemented') 62 | -------------------------------------------------------------------------------- /config/ext/keypoint_rcnn-backbone_ext_resnet50-b3ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: False 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | model: 24 | name: &model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &backbone_name 'custom_resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | layer1: 31 | name: 'Bottleneck4LargeResNet' 32 | bottleneck_channel: &bch 3 33 | ext_config: 34 | backbone_frozen: True 35 | threshold: 0.01 36 | ckpt: !join ['./resource/ckpt/ext/', *dataset_name, '-', *model_name, '-backbone_ext_', *backbone_name, '-b', *bch, 'ch.pt'] 37 | bottleneck_transformer: 38 | order: ['quantizer', 'dequantizer'] 39 | components: 40 | quantizer: 41 | params: 42 | num_bits: 8 43 | dequantizer: 44 | params: 45 | num_bits: 8 46 | params: 47 | num_classes: 2 48 | num_keypoints: 17 49 | pretrained: True 50 | experiment: &experiment !join [*dataset_name, '-', *model_name, '-backbone_', *backbone_name, '_from_keypoint_rcnn-backbone_resnet50-b', *bch, 'ch'] 51 | ckpt: !join ['./resource/ckpt/ghnd/', *experiment, '.pt'] 52 | 53 | train: 54 | num_epochs: 30 55 | batch_size: 2 56 | log_freq: 10000 57 | optimizer: 58 | type: 'SGD' 59 | params: 60 | lr: 0.001 61 | momentum: 0.9 62 | weight_decay: 0.0001 63 | scheduler: 64 | type: 'MultiStepLR' 65 | params: 66 | milestones: [15, 25] 67 | gamma: 0.1 68 | 69 | test: 70 | batch_size: 1 71 | -------------------------------------------------------------------------------- /src/visualizer.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | import torch 5 | from PIL import Image 6 | 7 | from models import get_model 8 | from myutils.common import file_util, yaml_util 9 | from utils import main_util, visual_util 10 | 11 | 12 | def get_argparser(): 13 | argparser = argparse.ArgumentParser(description=__doc__) 14 | argparser.add_argument('--config', required=True, help='yaml config file') 15 | argparser.add_argument('--device', default='cuda', help='device') 16 | argparser.add_argument('--json', help='dictionary to overwrite config') 17 | argparser.add_argument('--image', default=True, nargs='+', help='image dir/file paths for visualization') 18 | argparser.add_argument('--output', required=True, help='output dir path') 19 | return argparser 20 | 21 | 22 | def predict(model, input_file_path, device, output_dir_path): 23 | result, output, top_pred = visual_util.predict(model, Image.open(input_file_path).convert('RGB'), device) 24 | output_file_path = os.path.join(output_dir_path, os.path.basename(input_file_path)) 25 | vis_img = Image.fromarray(result[:, :, [2, 1, 0]]) 26 | vis_img.save(output_file_path) 27 | 28 | 29 | def visualize_predictions(model, input_file_paths, device, output_dir_path): 30 | model.eval() 31 | file_util.make_dirs(output_dir_path) 32 | for input_file_path in input_file_paths: 33 | if not file_util.check_if_exists(input_file_path): 34 | print('`{}` is not found.'.format(input_file_path)) 35 | continue 36 | 37 | if os.path.isfile(input_file_path): 38 | predict(model, input_file_path, device, output_dir_path) 39 | else: 40 | for sub_input_file_path in file_util.get_file_path_list(input_file_path, is_recursive=True): 41 | predict(model, sub_input_file_path, device, output_dir_path) 42 | 43 | 44 | def main(args): 45 | config = yaml_util.load_yaml_file(args.config) 46 | if args.json is not None: 47 | main_util.overwrite_config(config, args.json) 48 | 49 | device = torch.device(args.device) 50 | print(args) 51 | print('Creating model') 52 | model_config = config['model'] if 'model' in config else config.get('student_model', None) 53 | if model_config is None: 54 | raise ValueError('`{}` should contain model or student config at root'.format(args.config)) 55 | 56 | model = get_model(model_config, device, strict=False) 57 | visualize_predictions(model, args.image, device, args.output) 58 | 59 | 60 | if __name__ == '__main__': 61 | parser = get_argparser() 62 | main(parser.parse_args()) 63 | -------------------------------------------------------------------------------- /src/utils/data_util.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from structure.sampler import GroupedBatchSampler, create_aspect_ratio_groups 4 | from structure.transformer import ToTensor, RandomHorizontalFlip, Compose 5 | from utils import misc_util 6 | from utils.coco_util import get_coco 7 | 8 | 9 | def get_coco_dataset(split_dict, is_train): 10 | transforms = [ToTensor()] 11 | if is_train: 12 | transforms.append(RandomHorizontalFlip(0.5)) 13 | return get_coco(img_dir_path=split_dict['images'], ann_file_path=split_dict['annotations'], 14 | transforms=Compose(transforms), remove_non_annotated_imgs=split_dict['remove_non_annotated_imgs'], 15 | jpeg_quality=split_dict['jpeg_quality']) 16 | 17 | 18 | def get_coco_data_loaders(dataset_config, batch_size, distributed): 19 | num_workers = dataset_config['num_workers'] 20 | aspect_ratio_group_factor = dataset_config['aspect_ratio_group_factor'] 21 | dataset_splits = dataset_config['splits'] 22 | train_dataset = get_coco_dataset(dataset_splits['train'], True) 23 | val_dataset = get_coco_dataset(dataset_splits['val'], False) 24 | test_dataset = get_coco_dataset(dataset_splits['test'], False) 25 | 26 | print('Creating data loaders') 27 | if distributed: 28 | train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) 29 | val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset) 30 | test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset) 31 | else: 32 | train_sampler = torch.utils.data.RandomSampler(train_dataset) 33 | val_sampler = torch.utils.data.SequentialSampler(val_dataset) 34 | test_sampler = torch.utils.data.SequentialSampler(test_dataset) 35 | 36 | if aspect_ratio_group_factor >= 0: 37 | group_ids = create_aspect_ratio_groups(train_dataset, k=aspect_ratio_group_factor) 38 | train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, batch_size) 39 | else: 40 | train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size, drop_last=True) 41 | 42 | train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=train_batch_sampler, 43 | num_workers=num_workers, collate_fn=misc_util.collate_fn) 44 | val_data_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, sampler=val_sampler, 45 | num_workers=num_workers, collate_fn=misc_util.collate_fn) 46 | test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, sampler=test_sampler, 47 | num_workers=num_workers, collate_fn=misc_util.collate_fn) 48 | return train_sampler, train_data_loader, val_data_loader, test_data_loader 49 | -------------------------------------------------------------------------------- /src/models/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torchvision.models import detection 4 | 5 | from models.org import rcnn 6 | from myutils.common import file_util 7 | from utils import misc_util 8 | from structure.transformer import get_bottleneck_transformer 9 | 10 | 11 | def save_ckpt(model, optimizer, lr_scheduler, best_value, config, args, output_file_path): 12 | file_util.make_parent_dirs(output_file_path) 13 | model_state_dict =\ 14 | model.module.state_dict() if isinstance(model, nn.parallel.DistributedDataParallel) else model.state_dict() 15 | misc_util.save_on_master({'model': model_state_dict, 'optimizer': optimizer.state_dict(), 'best_value': best_value, 16 | 'lr_scheduler': lr_scheduler.state_dict(), 'config': config, 'args': args}, 17 | output_file_path) 18 | 19 | 20 | def load_ckpt(ckpt_file_path, model=None, optimizer=None, lr_scheduler=None, strict=True): 21 | if not file_util.check_if_exists(ckpt_file_path): 22 | print('ckpt file is not found at `{}`'.format(ckpt_file_path)) 23 | return None, None 24 | 25 | ckpt = torch.load(ckpt_file_path, map_location='cpu') 26 | if model is not None: 27 | print('Loading model parameters') 28 | model.load_state_dict(ckpt['model'], strict=strict) 29 | if optimizer is not None: 30 | print('Loading optimizer parameters') 31 | optimizer.load_state_dict(ckpt['optimizer']) 32 | if lr_scheduler is not None: 33 | print('Loading scheduler parameters') 34 | lr_scheduler.load_state_dict(ckpt['lr_scheduler']) 35 | return ckpt.get('best_value', 0.0), ckpt['config'], ckpt['args'] 36 | 37 | 38 | def get_model(model_config, device, strict=True, bottleneck_transformer=None): 39 | model_name = model_config['name'] 40 | ckpt_file_path = model_config['ckpt'] 41 | model_params_config = model_config['params'] 42 | if model_name in rcnn.MODEL_CLASS_DICT: 43 | backbone_config = model_config['backbone'] 44 | if bottleneck_transformer is None and 'bottleneck_transformer' in model_config: 45 | bottleneck_transformer = get_bottleneck_transformer(model_config['bottleneck_transformer']) 46 | 47 | model = rcnn.get_model(model_name, backbone_config=backbone_config, strict=strict, 48 | bottleneck_transformer=bottleneck_transformer, **model_params_config) 49 | if 'ext_config' in backbone_config: 50 | ext_config = backbone_config['ext_config'] 51 | load_ckpt(ext_config['ckpt'], model=model.backbone.body.get_ext_classifier()) 52 | strict = False 53 | else: 54 | raise ValueError('model_name `{}` is not expected'.format(model_name)) 55 | 56 | load_ckpt(ckpt_file_path, model=model, strict=strict) 57 | return model.to(device) 58 | 59 | 60 | def get_iou_types(model): 61 | model_without_ddp = model 62 | if isinstance(model, nn.parallel.DistributedDataParallel): 63 | model_without_ddp = model.module 64 | 65 | iou_type_list = ['bbox'] 66 | if isinstance(model_without_ddp, (detection.MaskRCNN, rcnn.MaskRCNN)): 67 | iou_type_list.append('segm') 68 | if isinstance(model_without_ddp, (detection.KeypointRCNN, rcnn.KeypointRCNN)): 69 | iou_type_list.append('keypoints') 70 | return iou_type_list 71 | -------------------------------------------------------------------------------- /config/hnd/mask_rcnn-backbone_resnet50-b12ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 12 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/mask_rcnn-backbone_resnet50-b15ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 15 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/mask_rcnn-backbone_resnet50-b3ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 3 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/mask_rcnn-backbone_resnet50-b6ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 6 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/mask_rcnn-backbone_resnet50-b9ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 9 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/faster_rcnn-backbone_resnet50-b12ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 12 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/faster_rcnn-backbone_resnet50-b15ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 15 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/faster_rcnn-backbone_resnet50-b3ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 3 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/faster_rcnn-backbone_resnet50-b6ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 6 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/faster_rcnn-backbone_resnet50-b9ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 9 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | scheduler: 84 | type: 'MultiStepLR' 85 | params: 86 | milestones: [5, 15] 87 | gamma: 0.1 88 | 89 | test: 90 | batch_size: 1 91 | -------------------------------------------------------------------------------- /config/hnd/keypoint_rcnn-backbone_resnet50-b3ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 3 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | scheduler: 86 | type: 'MultiStepLR' 87 | params: 88 | milestones: [9, 27] 89 | gamma: 0.1 90 | 91 | test: 92 | batch_size: 1 93 | -------------------------------------------------------------------------------- /config/hnd/keypoint_rcnn-backbone_resnet50-b6ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 6 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | scheduler: 86 | type: 'MultiStepLR' 87 | params: 88 | milestones: [9, 27] 89 | gamma: 0.1 90 | 91 | test: 92 | batch_size: 1 93 | -------------------------------------------------------------------------------- /config/hnd/keypoint_rcnn-backbone_resnet50-b9ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 9 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | scheduler: 86 | type: 'MultiStepLR' 87 | params: 88 | milestones: [9, 27] 89 | gamma: 0.1 90 | 91 | test: 92 | batch_size: 1 93 | -------------------------------------------------------------------------------- /config/hnd/keypoint_rcnn-backbone_resnet50-b12ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 12 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | scheduler: 86 | type: 'MultiStepLR' 87 | params: 88 | milestones: [9, 27] 89 | gamma: 0.1 90 | 91 | test: 92 | batch_size: 1 93 | -------------------------------------------------------------------------------- /config/hnd/keypoint_rcnn-backbone_resnet50-b15ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 15 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/hnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | scheduler: 86 | type: 'MultiStepLR' 87 | params: 88 | milestones: [9, 27] 89 | gamma: 0.1 90 | 91 | test: 92 | batch_size: 1 93 | -------------------------------------------------------------------------------- /src/distillation/tool.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | from torch import nn 4 | from torch.nn import DataParallel 5 | from torch.nn.parallel.distributed import DistributedDataParallel 6 | 7 | from distillation.loss import get_loss 8 | from models.org.rcnn import KeypointRCNN 9 | from myutils.pytorch import module_util 10 | 11 | 12 | class DistillationBox(nn.Module): 13 | def __init__(self, teacher_model, student_model, criterion_config): 14 | super().__init__() 15 | self.teacher_model = teacher_model 16 | self.student_model = student_model 17 | self.target_module_pairs = list() 18 | 19 | def extract_output(self, input, output): 20 | self.__dict__['distillation_box']['output'] = output 21 | 22 | teacher_model_without_dp = teacher_model.module if isinstance(teacher_model, DataParallel) else teacher_model 23 | student_model_without_ddp = \ 24 | student_model.module if isinstance(student_model, DistributedDataParallel) else student_model 25 | for loss_name, loss_config in criterion_config['terms'].items(): 26 | teacher_path, student_path = loss_config['ts_modules'] 27 | self.target_module_pairs.append((teacher_path, student_path)) 28 | teacher_module = module_util.get_module(teacher_model_without_dp, teacher_path) 29 | student_module = module_util.get_module(student_model_without_ddp, student_path) 30 | teacher_module.__dict__['distillation_box'] = {'loss_name': loss_name, 'path_from_root': teacher_path, 31 | 'is_teacher': True} 32 | student_module.__dict__['distillation_box'] = {'loss_name': loss_name, 'path_from_root': student_path, 33 | 'is_teacher': False} 34 | teacher_module.register_forward_hook(extract_output) 35 | student_module.register_forward_hook(extract_output) 36 | 37 | self.criterion = get_loss(criterion_config) 38 | self.require_adjustment = isinstance(student_model_without_ddp, KeypointRCNN) 39 | 40 | def forward(self, images, targets): 41 | teacher_model_without_dp =\ 42 | self.teacher_model.module if isinstance(self.teacher_model, DataParallel) else self.teacher_model 43 | student_model_without_ddp = \ 44 | self.student_model.module if isinstance(self.student_model, DistributedDataParallel) else self.student_model 45 | if self.require_adjustment: 46 | fixed_sizes = [random.choice(teacher_model_without_dp.transform.min_size) for _ in images] 47 | self.teacher_model(images, fixed_sizes=fixed_sizes) 48 | org_loss_dict = self.student_model(images, targets, fixed_sizes=fixed_sizes) 49 | else: 50 | self.teacher_model(images) 51 | org_loss_dict = self.student_model(images, targets) 52 | 53 | output_dict = dict() 54 | for teacher_path, student_path in self.target_module_pairs: 55 | teacher_dict = module_util.get_module(teacher_model_without_dp, teacher_path).__dict__['distillation_box'] 56 | student_dict = module_util.get_module(student_model_without_ddp, student_path).__dict__['distillation_box'] 57 | output_dict[teacher_dict['loss_name']] = ((teacher_dict['path_from_root'], teacher_dict['output']), 58 | (student_dict['path_from_root'], student_dict['output'])) 59 | 60 | total_loss = self.criterion(output_dict, org_loss_dict) 61 | return total_loss 62 | -------------------------------------------------------------------------------- /src/utils/main_util.py: -------------------------------------------------------------------------------- 1 | import builtins as __builtin__ 2 | import json 3 | import os 4 | import time 5 | 6 | import torch 7 | 8 | from models import get_iou_types 9 | from utils import misc_util 10 | from utils.coco_eval_util import CocoEvaluator 11 | from utils.coco_util import get_coco_api_from_dataset 12 | 13 | 14 | def overwrite_dict(org_dict, sub_dict): 15 | for sub_key, sub_value in sub_dict.items(): 16 | if sub_key in org_dict: 17 | if isinstance(sub_value, dict): 18 | overwrite_dict(org_dict[sub_key], sub_value) 19 | else: 20 | org_dict[sub_key] = sub_value 21 | else: 22 | org_dict[sub_key] = sub_value 23 | 24 | 25 | def overwrite_config(config, json_str): 26 | overwrite_dict(config, json.loads(json_str)) 27 | 28 | 29 | def setup_for_distributed(is_master): 30 | """ 31 | This function disables printing when not in master process 32 | """ 33 | builtin_print = __builtin__.print 34 | 35 | def print(*args, **kwargs): 36 | force = kwargs.pop('force', False) 37 | if is_master or force: 38 | builtin_print(*args, **kwargs) 39 | 40 | __builtin__.print = print 41 | 42 | 43 | def init_distributed_mode(world_size=1, dist_url='env://'): 44 | if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: 45 | rank = int(os.environ['RANK']) 46 | world_size = int(os.environ['WORLD_SIZE']) 47 | device_id = int(os.environ['LOCAL_RANK']) 48 | elif 'SLURM_PROCID' in os.environ: 49 | rank = int(os.environ['SLURM_PROCID']) 50 | device_id = rank % torch.cuda.device_count() 51 | else: 52 | print('Not using distributed mode') 53 | return False, None 54 | 55 | torch.cuda.set_device(device_id) 56 | dist_backend = 'nccl' 57 | print('| distributed init (rank {}): {}'.format(rank, dist_url), flush=True) 58 | torch.distributed.init_process_group(backend=dist_backend, init_method=dist_url, 59 | world_size=world_size, rank=rank) 60 | torch.distributed.barrier() 61 | setup_for_distributed(rank == 0) 62 | return True, [device_id] 63 | 64 | 65 | def warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor): 66 | def f(x): 67 | if x >= warmup_iters: 68 | return 1 69 | alpha = float(x) / warmup_iters 70 | return warmup_factor * (1 - alpha) + alpha 71 | 72 | return torch.optim.lr_scheduler.LambdaLR(optimizer, f) 73 | 74 | 75 | @torch.no_grad() 76 | def evaluate(model, data_loader, device): 77 | n_threads = torch.get_num_threads() 78 | # FIXME remove this and make paste_masks_in_image run on the GPU 79 | torch.set_num_threads(1) 80 | cpu_device = torch.device('cpu') 81 | model.eval() 82 | metric_logger = misc_util.MetricLogger(delimiter=' ') 83 | header = 'Test:' 84 | coco = get_coco_api_from_dataset(data_loader.dataset) 85 | iou_types = get_iou_types(model) 86 | coco_evaluator = CocoEvaluator(coco, iou_types) 87 | for image, targets in metric_logger.log_every(data_loader, 100, header): 88 | image = list(img.to(device) for img in image) 89 | targets = [{k: v.to(device) for k, v in t.items()} for t in targets] 90 | 91 | torch.cuda.synchronize() 92 | model_time = time.time() 93 | outputs = model(image) 94 | 95 | outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs] 96 | model_time = time.time() - model_time 97 | 98 | res = {target['image_id'].item(): output for target, output in zip(targets, outputs)} 99 | evaluator_time = time.time() 100 | coco_evaluator.update(res) 101 | evaluator_time = time.time() - evaluator_time 102 | metric_logger.update(model_time=model_time, evaluator_time=evaluator_time) 103 | 104 | # gather the stats from all processes 105 | metric_logger.synchronize_between_processes() 106 | print('Averaged stats:', metric_logger) 107 | coco_evaluator.synchronize_between_processes() 108 | 109 | # accumulate predictions from all images 110 | coco_evaluator.accumulate() 111 | coco_evaluator.summarize() 112 | torch.set_num_threads(n_threads) 113 | return coco_evaluator 114 | -------------------------------------------------------------------------------- /config/ghnd/mask_rcnn-backbone_resnet50-b12ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 12 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/mask_rcnn-backbone_resnet50-b3ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 3 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/mask_rcnn-backbone_resnet50-b6ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 6 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/mask_rcnn-backbone_resnet50-b9ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'mask_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'mask_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 9 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/faster_rcnn-backbone_resnet50-b12ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 12 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/faster_rcnn-backbone_resnet50-b3ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 3 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/faster_rcnn-backbone_resnet50-b6ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 6 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/faster_rcnn-backbone_resnet50-b9ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/instances_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/instances_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'faster_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 91 32 | pretrained: True 33 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 34 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 35 | 36 | student_model: 37 | name: &student_model_name 'faster_rcnn' 38 | backbone: 39 | name: &student_backbone_name 'custom_resnet50' 40 | params: 41 | pretrained: True 42 | freeze_layers: False 43 | layer1: 44 | name: 'Bottleneck4LargeResNet' 45 | bottleneck_channel: &bch 9 46 | bottleneck_transformer: 47 | order: ['quantizer', 'dequantizer'] 48 | components: 49 | quantizer: 50 | params: 51 | num_bits: 8 52 | dequantizer: 53 | params: 54 | num_bits: 8 55 | params: 56 | num_classes: 91 57 | pretrained: True 58 | distill_backbone_only: True 59 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 60 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 61 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 62 | 63 | train: 64 | num_epochs: 20 65 | batch_size: 4 66 | log_freq: 1000 67 | optimizer: 68 | type: 'Adam' 69 | params: 70 | lr: 0.001 71 | criterion: 72 | type: 'general' 73 | params: 74 | org_loss_factor: 0.0 75 | terms: 76 | layer1: 77 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 78 | criterion: 79 | type: 'MSELoss' 80 | params: 81 | reduction: 'sum' 82 | factor: 1.0 83 | layer2: 84 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 85 | criterion: 86 | type: 'MSELoss' 87 | params: 88 | reduction: 'sum' 89 | factor: 1.0 90 | layer3: 91 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 92 | criterion: 93 | type: 'MSELoss' 94 | params: 95 | reduction: 'sum' 96 | factor: 1.0 97 | layer4: 98 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 99 | criterion: 100 | type: 'MSELoss' 101 | params: 102 | reduction: 'sum' 103 | factor: 1.0 104 | scheduler: 105 | type: 'MultiStepLR' 106 | params: 107 | milestones: [5, 15] 108 | gamma: 0.1 109 | 110 | test: 111 | batch_size: 1 112 | -------------------------------------------------------------------------------- /config/ghnd/keypoint_rcnn-backbone_resnet50-b12ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 12 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | layer2: 86 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 87 | criterion: 88 | type: 'MSELoss' 89 | params: 90 | reduction: 'sum' 91 | factor: 1.0 92 | layer3: 93 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 94 | criterion: 95 | type: 'MSELoss' 96 | params: 97 | reduction: 'sum' 98 | factor: 1.0 99 | layer4: 100 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 101 | criterion: 102 | type: 'MSELoss' 103 | params: 104 | reduction: 'sum' 105 | factor: 1.0 106 | scheduler: 107 | type: 'MultiStepLR' 108 | params: 109 | milestones: [9, 27] 110 | gamma: 0.1 111 | 112 | test: 113 | batch_size: 1 114 | -------------------------------------------------------------------------------- /config/ghnd/keypoint_rcnn-backbone_resnet50-b3ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 3 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | layer2: 86 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 87 | criterion: 88 | type: 'MSELoss' 89 | params: 90 | reduction: 'sum' 91 | factor: 1.0 92 | layer3: 93 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 94 | criterion: 95 | type: 'MSELoss' 96 | params: 97 | reduction: 'sum' 98 | factor: 1.0 99 | layer4: 100 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 101 | criterion: 102 | type: 'MSELoss' 103 | params: 104 | reduction: 'sum' 105 | factor: 1.0 106 | scheduler: 107 | type: 'MultiStepLR' 108 | params: 109 | milestones: [9, 27] 110 | gamma: 0.1 111 | 112 | test: 113 | batch_size: 1 114 | -------------------------------------------------------------------------------- /config/ghnd/keypoint_rcnn-backbone_resnet50-b6ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 6 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | layer2: 86 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 87 | criterion: 88 | type: 'MSELoss' 89 | params: 90 | reduction: 'sum' 91 | factor: 1.0 92 | layer3: 93 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 94 | criterion: 95 | type: 'MSELoss' 96 | params: 97 | reduction: 'sum' 98 | factor: 1.0 99 | layer4: 100 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 101 | criterion: 102 | type: 'MSELoss' 103 | params: 104 | reduction: 'sum' 105 | factor: 1.0 106 | scheduler: 107 | type: 'MultiStepLR' 108 | params: 109 | milestones: [9, 27] 110 | gamma: 0.1 111 | 112 | test: 113 | batch_size: 1 114 | -------------------------------------------------------------------------------- /config/ghnd/keypoint_rcnn-backbone_resnet50-b9ch.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | name: &dataset_name 'coco2017' 3 | root: &root_dir !join ['./resource/dataset/', *dataset_name] 4 | num_workers: 4 5 | aspect_ratio_group_factor: 3 6 | splits: 7 | train: 8 | images: !join [*root_dir, '/train2017'] 9 | annotations: !join [*root_dir, '/annotations/person_keypoints_train2017.json'] 10 | remove_non_annotated_imgs: True 11 | jpeg_quality: 12 | val: 13 | images: !join [*root_dir, '/val2017'] 14 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 15 | remove_non_annotated_imgs: False 16 | jpeg_quality: 17 | test: 18 | images: !join [*root_dir, '/val2017'] 19 | annotations: !join [*root_dir, '/annotations/person_keypoints_val2017.json'] 20 | remove_non_annotated_imgs: False 21 | jpeg_quality: 22 | 23 | teacher_model: 24 | name: &teacher_model_name 'keypoint_rcnn' 25 | backbone: 26 | name: &teacher_backbone_name 'resnet50' 27 | params: 28 | pretrained: True 29 | freeze_layers: True 30 | params: 31 | num_classes: 2 32 | num_keypoints: 17 33 | pretrained: True 34 | experiment: &teacher_experiment !join [*dataset_name, '-', *teacher_model_name, '-backbone_', *teacher_backbone_name] 35 | ckpt: !join ['./resource/ckpt/org/', *teacher_experiment, '.pt'] 36 | 37 | student_model: 38 | name: &student_model_name 'keypoint_rcnn' 39 | backbone: 40 | name: &student_backbone_name 'custom_resnet50' 41 | params: 42 | pretrained: True 43 | freeze_layers: False 44 | layer1: 45 | name: 'Bottleneck4LargeResNet' 46 | bottleneck_channel: &bch 9 47 | bottleneck_transformer: 48 | order: ['quantizer', 'dequantizer'] 49 | components: 50 | quantizer: 51 | params: 52 | num_bits: 8 53 | dequantizer: 54 | params: 55 | num_bits: 8 56 | params: 57 | num_classes: 2 58 | num_keypoints: 17 59 | pretrained: True 60 | distill_backbone_only: True 61 | frozen_modules: ['backbone.body.layer2', 'backbone.body.layer3', 'backbone.body.layer4', 'backbone.fpn', 'rpn', 'roi_heads'] 62 | experiment: &student_experiment !join [*dataset_name, '-', *student_model_name, '-backbone_', *student_backbone_name, '_from_', *teacher_model_name, '-backbone_', *teacher_backbone_name, '-b', *bch, 'ch'] 63 | ckpt: !join ['./resource/ckpt/ghnd/', *student_experiment, '.pt'] 64 | 65 | train: 66 | num_epochs: 35 67 | batch_size: 4 68 | log_freq: 1000 69 | optimizer: 70 | type: 'Adam' 71 | params: 72 | lr: 0.001 73 | criterion: 74 | type: 'general' 75 | params: 76 | org_loss_factor: 0.0 77 | terms: 78 | layer1: 79 | ts_modules: ['backbone.body.layer1', 'backbone.body.layer1'] 80 | criterion: 81 | type: 'MSELoss' 82 | params: 83 | reduction: 'sum' 84 | factor: 1.0 85 | layer2: 86 | ts_modules: ['backbone.body.layer2', 'backbone.body.layer2'] 87 | criterion: 88 | type: 'MSELoss' 89 | params: 90 | reduction: 'sum' 91 | factor: 1.0 92 | layer3: 93 | ts_modules: ['backbone.body.layer3', 'backbone.body.layer3'] 94 | criterion: 95 | type: 'MSELoss' 96 | params: 97 | reduction: 'sum' 98 | factor: 1.0 99 | layer4: 100 | ts_modules: ['backbone.body.layer4', 'backbone.body.layer4'] 101 | criterion: 102 | type: 'MSELoss' 103 | params: 104 | reduction: 'sum' 105 | factor: 1.0 106 | scheduler: 107 | type: 'MultiStepLR' 108 | params: 109 | milestones: [9, 27] 110 | gamma: 0.1 111 | 112 | test: 113 | batch_size: 1 114 | -------------------------------------------------------------------------------- /src/models/mimic/resnet_layer.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from models.ext.classifier import Ext4ResNet 4 | from models.mimic.base import BottleneckBase4Ext, ExtEncoder 5 | 6 | 7 | class Bottleneck4SmallResNet(BottleneckBase4Ext): 8 | def __init__(self, bottleneck_channel, ext_config, bottleneck_transformer): 9 | encoder = nn.Sequential( 10 | nn.Conv2d(64, 64, kernel_size=2, padding=1, bias=False), 11 | nn.BatchNorm2d(64), 12 | nn.Conv2d(64, 256, kernel_size=2, padding=1, bias=False), 13 | nn.BatchNorm2d(256), 14 | nn.ReLU(inplace=True), 15 | nn.Conv2d(256, 64, kernel_size=2, padding=1, bias=False), 16 | nn.BatchNorm2d(64), 17 | nn.Conv2d(64, bottleneck_channel, kernel_size=2, padding=1, bias=False) 18 | ) 19 | decoder = nn.Sequential( 20 | nn.BatchNorm2d(bottleneck_channel), 21 | nn.ReLU(inplace=True), 22 | nn.Conv2d(bottleneck_channel, 64, kernel_size=2, bias=False), 23 | nn.BatchNorm2d(64), 24 | nn.Conv2d(64, 128, kernel_size=2, bias=False), 25 | nn.BatchNorm2d(128), 26 | nn.ReLU(inplace=True), 27 | nn.Conv2d(128, 64, kernel_size=2, bias=False), 28 | nn.BatchNorm2d(64), 29 | nn.Conv2d(64, 64, kernel_size=2, bias=False), 30 | nn.BatchNorm2d(64), 31 | nn.ReLU(inplace=True) 32 | ) 33 | encoder = ExtEncoder(encoder, Ext4ResNet(64) if ext_config is not None else None, ext_config) 34 | super().__init__(encoder=encoder, decoder=decoder, bottleneck_transformer=bottleneck_transformer) 35 | 36 | def get_ext_classifier(self): 37 | return self.encoder.get_ext_classifier() 38 | 39 | 40 | class Bottleneck4LargeResNet(BottleneckBase4Ext): 41 | def __init__(self, bottleneck_channel, ext_config, bottleneck_transformer): 42 | encoder = nn.Sequential( 43 | nn.Conv2d(64, 64, kernel_size=2, padding=1, bias=False), 44 | nn.BatchNorm2d(64), 45 | nn.Conv2d(64, 256, kernel_size=2, padding=1, bias=False), 46 | nn.BatchNorm2d(256), 47 | nn.ReLU(inplace=True), 48 | nn.Conv2d(256, 64, kernel_size=2, padding=1, bias=False), 49 | nn.BatchNorm2d(64), 50 | nn.Conv2d(64, bottleneck_channel, kernel_size=2, padding=1, bias=False) 51 | ) 52 | decoder = nn.Sequential( 53 | nn.BatchNorm2d(bottleneck_channel), 54 | nn.ReLU(inplace=True), 55 | nn.Conv2d(bottleneck_channel, 64, kernel_size=2, bias=False), 56 | nn.BatchNorm2d(64), 57 | nn.Conv2d(64, 128, kernel_size=2, bias=False), 58 | nn.BatchNorm2d(128), 59 | nn.ReLU(inplace=True), 60 | nn.Conv2d(128, 256, kernel_size=2, bias=False), 61 | nn.BatchNorm2d(256), 62 | nn.Conv2d(256, 256, kernel_size=2, bias=False), 63 | nn.BatchNorm2d(256), 64 | nn.ReLU(inplace=True) 65 | ) 66 | encoder = ExtEncoder(encoder, Ext4ResNet(64) if ext_config is not None else None, ext_config) 67 | super().__init__(encoder=encoder, decoder=decoder, bottleneck_transformer=bottleneck_transformer) 68 | 69 | def get_ext_classifier(self): 70 | return self.encoder.get_ext_classifier() 71 | 72 | 73 | def get_mimic_layers(backbone_name, backbone_config, bottleneck_transformer=None): 74 | layer1, layer2, layer3, layer4 = None, None, None, None 75 | backbone_params_config = backbone_config['params'] 76 | layer1_config = backbone_params_config.get('layer1', None) 77 | if layer1_config is not None: 78 | layer1_name = layer1_config['name'] 79 | ext_config = backbone_config.get('ext_config', None) 80 | if layer1_name == 'Bottleneck4SmallResNet' and backbone_name in {'custom_resnet18', 'custom_resnet34'}: 81 | layer1 = Bottleneck4LargeResNet(layer1_config['bottleneck_channel'], ext_config, bottleneck_transformer) 82 | elif layer1_name == 'Bottleneck4LargeResNet'\ 83 | and backbone_name in {'custom_resnet50', 'custom_resnet101', 'custom_resnet152'}: 84 | layer1 = Bottleneck4LargeResNet(layer1_config['bottleneck_channel'], ext_config, bottleneck_transformer) 85 | else: 86 | raise ValueError('layer1_name `{}` is not expected'.format(layer1_name)) 87 | return layer1, layer2, layer3, layer4 88 | -------------------------------------------------------------------------------- /src/models/ext/backbone.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | 3 | from torch import nn 4 | from torch.jit.annotations import Dict 5 | from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork, LastLevelMaxPool 6 | 7 | from models.mimic.base import BottleneckBase4Ext, ExtEncoder 8 | from myutils.pytorch import module_util 9 | 10 | 11 | def has_only_empty_bbox(target): 12 | return all(any(o <= 1 for o in box[2:]) for box in target['boxes']) 13 | 14 | 15 | def count_visible_keypoints(target): 16 | return sum(sum(1 for row in kp if row[2] > 0) for kp in target['keypoints']) 17 | 18 | 19 | def check_if_valid_target(target, min_keypoints_per_image=10): 20 | # if it's empty, there is no annotation 21 | if len(target) == 0: 22 | return False 23 | # if all boxes have close to zero area, there is no annotation 24 | if has_only_empty_bbox(target): 25 | return False 26 | # keypoints task have a slight different criteria for considering 27 | # if an annotation is valid 28 | if 'keypoints' not in target: 29 | return True 30 | # for keypoint detection tasks, only consider valid images those 31 | # containing at least min_keypoints_per_image 32 | if count_visible_keypoints(target) >= min_keypoints_per_image: 33 | return True 34 | return False 35 | 36 | 37 | def check_if_includes_ext(module): 38 | return isinstance(module, BottleneckBase4Ext) and isinstance(module.encoder, ExtEncoder) 39 | 40 | 41 | class ExtIntermediateLayerGetter(nn.ModuleDict): 42 | _version = 2 43 | __constants__ = ['layers'] 44 | __annotations__ = { 45 | "return_layers": Dict[str, str], 46 | } 47 | 48 | def __init__(self, model, return_layers, ext_config): 49 | if not set(return_layers).issubset([name for name, _ in model.named_children()]): 50 | raise ValueError("return_layers are not present in model") 51 | 52 | orig_return_layers = return_layers 53 | return_layers = {k: v for k, v in return_layers.items()} 54 | layers = OrderedDict() 55 | for name, module in model.named_children(): 56 | layers[name] = module 57 | if name in return_layers: 58 | del return_layers[name] 59 | if not return_layers: 60 | break 61 | 62 | super().__init__(layers) 63 | self.return_layers = orig_return_layers 64 | self.threshold = ext_config['threshold'] 65 | self.ext_training = False 66 | 67 | def get_ext_classifier(self): 68 | for _, module in self.items(): 69 | if check_if_includes_ext(module): 70 | return module.get_ext_classifier() 71 | return None 72 | 73 | def forward(self, x): 74 | out = OrderedDict() 75 | ext_x = None 76 | for name, module in self.items(): 77 | x = module(x) 78 | if name in self.return_layers: 79 | out_name = self.return_layers[name] 80 | out[out_name] = x 81 | if check_if_includes_ext(module): 82 | x, ext_x = x 83 | out[out_name] = x 84 | if x is None: 85 | return None, ext_x 86 | elif self.ext_training: 87 | return x, ext_x 88 | return out, ext_x 89 | 90 | 91 | class ExtBackboneWithFPN(nn.Module): 92 | def __init__(self, backbone, return_layers, in_channels_list, out_channels, ext_config): 93 | super().__init__() 94 | if ext_config.get('backbone_frozen', False): 95 | module_util.freeze_module_params(backbone) 96 | 97 | self.body = ExtIntermediateLayerGetter(backbone, return_layers=return_layers, ext_config=ext_config) 98 | self.fpn = FeaturePyramidNetwork( 99 | in_channels_list=in_channels_list, 100 | out_channels=out_channels, 101 | extra_blocks=LastLevelMaxPool(), 102 | ) 103 | self.out_channels = out_channels 104 | self.split = False 105 | 106 | def forward(self, x): 107 | if self.split: 108 | z = self.body(x) 109 | return self.fpn(z) 110 | 111 | z, ext_z = self.body(x) 112 | if (not self.training and z is None) or self.body.ext_training: 113 | return None, ext_z 114 | elif self.training: 115 | return z, ext_z 116 | return self.fpn(z), ext_z 117 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # HND & GHND for Object Detectors 2 | Head Network Distillation (HND) and Generalized HND for Faster, Mask, and Keypoint R-CNNs 3 | - "Neural Compression and Filtering for Edge-assisted Real-time Object Detection in Challenged Networks," [ICPR 2020](https://www.micc.unifi.it/icpr2020/) 4 | [[PDF](https://ieeexplore.ieee.org/document/9412388)] [[Supp](https://ieeexplore.ieee.org/document/9412388/media#media)] [[Preprint](https://arxiv.org/abs/2007.15818)] 5 | - "Split Computing for Complex Object Detectors: Challenges and Preliminary Results," [ACM MobiCom 2020 Workshop EMDL '20](https://emdl20.github.io/index.html) 6 | [[PDF (Open Access)](https://dl.acm.org/doi/abs/10.1145/3410338.3412338)] [[Preprint](https://arxiv.org/abs/2007.13312)] 7 | 8 | Looking for **image classification tasks**? Take a look at https://github.com/yoshitomo-matsubara/head-network-distillation as well 9 | 10 | 11 | ![GHND and Neural Filter](img/ghnd_and_nf.png) 12 | 13 | ## Citations 14 | ```bibtex 15 | @inproceedings{matsubara2021neural, 16 | title={Neural Compression and Filtering for Edge-assisted Real-time Object Detection in Challenged Networks}, 17 | author={Yoshitomo Matsubara and Marco Levorato}, 18 | booktitle={2020 25th International Conference on Pattern Recognition (ICPR)}, 19 | pages={2272--2279}, 20 | year={2021} 21 | } 22 | 23 | @inproceedings{matsubara2020split, 24 | title={Split Computing for Complex Object Detectors: Challenges and Preliminary Results}, 25 | author={Matsubara, Yoshitomo and Levorato, Marco}, 26 | booktitle={Proceedings of the 4th International Workshop on Embedded and Mobile Deep Learning}, 27 | pages={7--12}, 28 | year={2020} 29 | } 30 | 31 | ``` 32 | 33 | ## Requirements 34 | - Python 3.6 35 | - pipenv 36 | - [myutils](https://github.com/yoshitomo-matsubara/myutils) 37 | 38 | ## How to clone 39 | ``` 40 | git clone https://github.com/yoshitomo-matsubara/hnd-ghnd-object-detectors.git 41 | cd hnd-ghnd-object-detectors/ 42 | git submodule init 43 | git submodule update --recursive --remote 44 | pipenv install 45 | ``` 46 | It is not necessary to use pipenv, and you can instead manually install the required packages listed in [Pipfile](Pipfile), using pip3. 47 | 48 | ## COCO 2017 Dataset 49 | ``` 50 | mkdir -p ./resource/dataset/coco2017 51 | cd ./resource/dataset/coco2017 52 | wget http://images.cocodataset.org/zips/train2017.zip 53 | wget http://images.cocodataset.org/zips/val2017.zip 54 | wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip 55 | unzip -q train2017.zip 56 | unzip -q val2017.zip 57 | unzip -q annotations_trainval2017.zip 58 | ``` 59 | 60 | ## Checkpoints with trained model weights 61 | 1. Download [emdl2020.zip](https://github.com/yoshitomo-matsubara/hnd-ghnd-object-detectors/releases/download/google_drive-to-github/emdl2020.zip) 62 | 2. Unzip **emdl2020.zip** at the root directory of this repository so that you can use the checkpoints with yaml config files under [config/hnd/](config/hnd/) 63 | 3. Download [icpr2020-faster_rcnn.zip](https://github.com/yoshitomo-matsubara/hnd-ghnd-object-detectors/releases/download/google_drive-to-github/icpr2020-faster_rcnn.zip), [icpr2020-mask_rcnn.zip](https://github.com/yoshitomo-matsubara/hnd-ghnd-object-detectors/releases/download/google_drive-to-github/icpr2020-mask_rcnn.zip), and [icpr2020-keypoint_rcnn.zip](https://github.com/yoshitomo-matsubara/hnd-ghnd-object-detectors/releases/download/google_drive-to-github/icpr2020-keypoint_rcnn.zip) 64 | 4. Unzip the three zip files at the root directory of this repository so that you can use the checkpoints with yaml config files under [config/hnd/](config/hnd/) and [config/ghnd/](config/ghnd/) 65 | 5. Test the trained models using the checkpoints and yaml config files 66 | e.g., Faster R-CNN with 3 output channels for bottleneck 67 | ``` 68 | pipenv run python src/mimic_runner.py --config config/hnd/faster_rcnn-backbone_resnet50-b3ch.yaml 69 | pipenv run python src/mimic_runner.py --config config/ghnd/faster_rcnn-backbone_resnet50-b3ch.yaml 70 | ``` 71 | 72 | ## Distilling head portion of R-CNNs 73 | If you have already downloaded our trained model weights above, you should move the ckpt files in `resource/ckpt/` to somewhere else or change ckpt file path (`ckpt` under `student_model`) in config files. 74 | 75 | ### Bottleneck-injected Faster R-CNN with ResNet-50 and FPN 76 | e.g., Bottleneck with 3 output channels 77 | ``` 78 | # HND 79 | pipenv run python src/mimic_runner.py --config config/hnd/faster_rcnn-backbone_resnet50-b3ch.yaml -distill 80 | 81 | # GHND 82 | pipenv run python src/mimic_runner.py --config config/ghnd/faster_rcnn-backbone_resnet50-b3ch.yaml -distill 83 | ``` 84 | 85 | ### Bottleneck-injected Mask R-CNN with ResNet-50 and FPN 86 | e.g., Bottleneck with 3 output channels 87 | ``` 88 | # HND 89 | pipenv run python src/mimic_runner.py --config config/hnd/mask_rcnn-backbone_resnet50-b3ch.yaml -distill 90 | 91 | # GHND 92 | pipenv run python src/mimic_runner.py --config config/ghnd/mask_rcnn-backbone_resnet50-b3ch.yaml -distill 93 | ``` 94 | 95 | ### Bottleneck-injected Keypoint R-CNN with ResNet-50 and FPN 96 | e.g., Bottleneck with 3 output channels 97 | ``` 98 | # HND 99 | pipenv run python src/mimic_runner.py --config config/hnd/keypoint_rcnn-backbone_resnet50-b3ch.yaml -distill 100 | 101 | # GHND 102 | pipenv run python src/mimic_runner.py --config config/ghnd/keypoint_rcnn-backbone_resnet50-b3ch.yaml -distill 103 | ``` 104 | 105 | ## Training a neural filter on top of our trained, bottleneck-injected Keypoint R-CNN 106 | ``` 107 | pipenv run python src/ext_runner.py --config config/ext/keypoint_rcnn-backbone_ext_resnet50-b3ch.yaml -train 108 | ``` 109 | 110 | ## References 111 | - [pytorch/vision/references/detection/](https://github.com/pytorch/vision/tree/master/references/detection) 112 | - [code for visualization in the object detection tutorial](https://github.com/pytorch/vision/issues/1610) 113 | -------------------------------------------------------------------------------- /src/coco_runner.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import datetime 3 | import math 4 | import sys 5 | import time 6 | 7 | import torch 8 | from torch import distributed as dist 9 | from torch import nn 10 | 11 | from models import get_model, load_ckpt, save_ckpt 12 | from myutils.common import file_util, yaml_util 13 | from myutils.pytorch import func_util 14 | from utils import data_util, main_util, misc_util 15 | 16 | 17 | def get_argparser(): 18 | argparser = argparse.ArgumentParser(description=__doc__) 19 | argparser.add_argument('--config', required=True, help='yaml config file') 20 | argparser.add_argument('--device', default='cuda', help='device') 21 | argparser.add_argument('--json', help='dictionary to overwrite config') 22 | argparser.add_argument('-train', action='store_true', help='train a model') 23 | # distributed training parameters 24 | argparser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') 25 | argparser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') 26 | return argparser 27 | 28 | 29 | def train_model(model, optimizer, data_loader, device, epoch, log_freq): 30 | model.train() 31 | metric_logger = misc_util.MetricLogger(delimiter=' ') 32 | metric_logger.add_meter('lr', misc_util.SmoothedValue(window_size=1, fmt='{value:.6f}')) 33 | header = 'Epoch: [{}]'.format(epoch) 34 | lr_scheduler = None 35 | if epoch == 0: 36 | warmup_factor = 1.0 / 1000.0 37 | warmup_iters = min(1000, len(data_loader) - 1) 38 | lr_scheduler = main_util.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor) 39 | 40 | for images, targets in metric_logger.log_every(data_loader, log_freq, header): 41 | images = list(image.to(device) for image in images) 42 | targets = [{k: v.to(device) for k, v in t.items()} for t in targets] 43 | loss_dict = model(images, targets) 44 | losses = sum(loss for loss in loss_dict.values()) 45 | 46 | # reduce losses over all GPUs for logging purposes 47 | loss_dict_reduced = misc_util.reduce_dict(loss_dict) 48 | losses_reduced = sum(loss for loss in loss_dict_reduced.values()) 49 | loss_value = losses_reduced.item() 50 | 51 | if not math.isfinite(loss_value): 52 | print('Loss is {}, stopping training'.format(loss_value)) 53 | print(loss_dict_reduced) 54 | sys.exit(1) 55 | 56 | optimizer.zero_grad() 57 | losses.backward() 58 | optimizer.step() 59 | if lr_scheduler is not None: 60 | lr_scheduler.step() 61 | 62 | metric_logger.update(loss=losses_reduced, **loss_dict_reduced) 63 | metric_logger.update(lr=optimizer.param_groups[0]['lr']) 64 | 65 | 66 | def train(model, train_sampler, train_data_loader, val_data_loader, device, distributed, config, args, ckpt_file_path): 67 | train_config = config['train'] 68 | optim_config = train_config['optimizer'] 69 | optimizer = func_util.get_optimizer(model, optim_config['type'], optim_config['params']) 70 | scheduler_config = train_config['scheduler'] 71 | lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) 72 | best_val_map = 0.0 73 | if file_util.check_if_exists(ckpt_file_path): 74 | best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) 75 | 76 | num_epochs = train_config['num_epochs'] 77 | log_freq = train_config['log_freq'] 78 | start_time = time.time() 79 | for epoch in range(num_epochs): 80 | if distributed: 81 | train_sampler.set_epoch(epoch) 82 | 83 | train_model(model, optimizer, train_data_loader, device, epoch, log_freq) 84 | lr_scheduler.step() 85 | 86 | # evaluate after every epoch 87 | coco_evaluator = main_util.evaluate(model, val_data_loader, device=device) 88 | # Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] 89 | val_map = coco_evaluator.coco_eval['bbox'].stats[0] 90 | if val_map > best_val_map: 91 | print('Updating ckpt (Best BBox mAP: {:.4f} -> {:.4f})'.format(best_val_map, val_map)) 92 | best_val_map = val_map 93 | save_ckpt(model, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path) 94 | lr_scheduler.step() 95 | 96 | dist.barrier() 97 | total_time = time.time() - start_time 98 | total_time_str = str(datetime.timedelta(seconds=int(total_time))) 99 | print('Training time {}'.format(total_time_str)) 100 | 101 | 102 | def main(args): 103 | distributed, device_ids = main_util.init_distributed_mode(args.world_size, args.dist_url) 104 | config = yaml_util.load_yaml_file(args.config) 105 | if args.json is not None: 106 | main_util.overwrite_config(config, args.json) 107 | 108 | device = torch.device(args.device) 109 | print(args) 110 | 111 | print('Loading data') 112 | train_config = config['train'] 113 | train_sampler, train_data_loader, val_data_loader, test_data_loader =\ 114 | data_util.get_coco_data_loaders(config['dataset'], train_config['batch_size'], distributed) 115 | 116 | print('Creating model') 117 | model_config = config['model'] 118 | model = get_model(model_config, device) 119 | if distributed: 120 | model = nn.parallel.DistributedDataParallel(model, device_ids=device_ids) 121 | 122 | if args.train: 123 | print('Start training') 124 | start_time = time.time() 125 | train(model, train_sampler, train_data_loader, val_data_loader, device, distributed, 126 | config, args, model_config['ckpt']) 127 | total_time = time.time() - start_time 128 | total_time_str = str(datetime.timedelta(seconds=int(total_time))) 129 | print('Training time {}'.format(total_time_str)) 130 | main_util.evaluate(model, test_data_loader, device=device) 131 | 132 | 133 | if __name__ == '__main__': 134 | parser = get_argparser() 135 | main(parser.parse_args()) 136 | -------------------------------------------------------------------------------- /src/structure/transformer.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | 4 | import torch 5 | from PIL import Image 6 | from torchvision.transforms import functional 7 | 8 | from myutils.common import file_util 9 | from myutils.pytorch import tensor_util 10 | 11 | 12 | def _flip_coco_person_keypoints(kps, width): 13 | flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] 14 | flipped_data = kps[:, flip_inds] 15 | flipped_data[..., 0] = width - flipped_data[..., 0] 16 | # Maintain COCO convention that if visibility == 0, then x, y = 0 17 | inds = flipped_data[..., 2] == 0 18 | flipped_data[inds] = 0 19 | return flipped_data 20 | 21 | 22 | class Compose(object): 23 | def __init__(self, transforms): 24 | self.transforms = transforms 25 | 26 | def __call__(self, image, target): 27 | for t in self.transforms: 28 | image, target = t(image, target) 29 | return image, target 30 | 31 | 32 | class RandomHorizontalFlip(object): 33 | def __init__(self, prob): 34 | self.prob = prob 35 | 36 | def __call__(self, image, target): 37 | if random.random() < self.prob: 38 | height, width = image.shape[-2:] 39 | image = image.flip(-1) 40 | bbox = target['boxes'] 41 | bbox[:, [0, 2]] = width - bbox[:, [2, 0]] 42 | target['boxes'] = bbox 43 | if 'masks' in target: 44 | target['masks'] = target['masks'].flip(-1) 45 | if 'keypoints' in target: 46 | keypoints = target['keypoints'] 47 | keypoints = _flip_coco_person_keypoints(keypoints, width) 48 | target['keypoints'] = keypoints 49 | return image, target 50 | 51 | 52 | class ToTensor(object): 53 | def __call__(self, image, target): 54 | image = functional.to_tensor(image) 55 | return image, target 56 | 57 | 58 | class DataLogger(object): 59 | def __init__(self, num_bits=8): 60 | self.num_bits4quant = num_bits 61 | self.data_size_list = list() 62 | self.fp16_data_size_list = list() 63 | self.quantized_data_size_list = list() 64 | self.tensor_shape_list = list() 65 | 66 | def get_data(self): 67 | return self.data_size_list.copy(), self.fp16_data_size_list,\ 68 | self.quantized_data_size_list.copy(), self.tensor_shape_list.copy() 69 | 70 | def clear(self): 71 | self.data_size_list.clear() 72 | self.fp16_data_size_list.clear() 73 | self.quantized_data_size_list.clear() 74 | self.tensor_shape_list.clear() 75 | 76 | def __call__(self, z, target): 77 | if z is None: 78 | data_size = 0.0 79 | fp16_data_size = 0.0 80 | quantized_data_size = 0.0 81 | else: 82 | data_size = file_util.get_binary_object_size(z) 83 | fp16_data_size = None if not isinstance(z, torch.Tensor) else file_util.get_binary_object_size(z.short()) 84 | quantized_data_size = None if not isinstance(z, torch.Tensor)\ 85 | else file_util.get_binary_object_size(tensor_util.quantize_tensor(z, num_bits=self.num_bits4quant)) 86 | 87 | self.data_size_list.append(data_size) 88 | self.fp16_data_size_list.append(fp16_data_size) 89 | self.quantized_data_size_list.append(quantized_data_size) 90 | self.tensor_shape_list.append([0, 0, 0] if z is None else [z.shape[1], z.shape[2], z.shape[3]]) 91 | return z, target 92 | 93 | 94 | class JpegCompressor(object): 95 | def __init__(self, jpeg_quality=95, tmp_dir_path='./tmp/'): 96 | self.jpeg_quality = jpeg_quality 97 | self.tmp_dir_path = tmp_dir_path 98 | file_util.make_dirs(tmp_dir_path) 99 | 100 | def save_image(self, z, output_file_path): 101 | qz = tensor_util.quantize_tensor(z) 102 | img = Image.fromarray(qz.tensor.permute(1, 2, 0).cpu().numpy()) 103 | img.save(output_file_path, format='jpeg', quality=self.jpeg_quality) 104 | return qz 105 | 106 | def __call__(self, z, target): 107 | if (z.dim() == 3 and z.shape[0] == 3) or (z.dim() == 4 and z.shape[0] == 1 and z.shape[1] == 3): 108 | if z.dim() == 4: 109 | z = z.squeeze(0) 110 | 111 | file_path = os.path.join(self.tmp_dir_path, '{}.jpg'.format(hash(z))) 112 | qz = self.save_image(z, file_path) 113 | return (file_path, qz), target 114 | return z, target 115 | 116 | 117 | class JpegDecompressor(object): 118 | def __init__(self, tmp_dir_path='./tmp/', target_dim=4): 119 | self.tmp_dir_path = tmp_dir_path 120 | self.target_dim = target_dim 121 | 122 | def __call__(self, z, target): 123 | if isinstance(z, tuple) and isinstance(z[0], str): 124 | img = Image.open(z[0]).convert('RGB') 125 | qz = z[1] 126 | img = qz.scale * (functional.to_tensor(img) * 255.0 - qz.zero_point) 127 | return img if self.target_dim != 4 else img.unsqueeze(0), target 128 | return z, target 129 | 130 | 131 | class Quantizer(object): 132 | def __init__(self, num_bits=8): 133 | self.num_bits = num_bits 134 | 135 | def __call__(self, z, target): 136 | if self.num_bits == 16: 137 | return z.half(), target 138 | 139 | qz = tensor_util.quantize_tensor(z, num_bits=self.num_bits) 140 | return qz, target 141 | 142 | 143 | class Dequantizer(object): 144 | def __init__(self, num_bits=8): 145 | # num_bits should be the same as Quantizer 146 | self.num_bits = num_bits 147 | 148 | def __call__(self, qz, target): 149 | if self.num_bits == 16: 150 | return qz.float(), target 151 | 152 | z = tensor_util.dequantize_tensor(qz) 153 | return z, target 154 | 155 | 156 | TRANSFORMER_CLASS_DICT = { 157 | 'jpeg_compressor': JpegCompressor, 158 | 'jpeg_decompressor': JpegDecompressor, 159 | 'quantizer': Quantizer, 160 | 'dequantizer': Dequantizer 161 | } 162 | 163 | 164 | def get_bottleneck_transformer(transformer_config): 165 | component_list = list() 166 | components_config = transformer_config['components'] 167 | for component_name in transformer_config['order']: 168 | param_config = components_config[component_name]['params'] 169 | if component_name not in TRANSFORMER_CLASS_DICT: 170 | raise KeyError('transformer `{}` is not expected'.format(component_name)) 171 | 172 | obj_class = TRANSFORMER_CLASS_DICT[component_name] 173 | component_list.append(obj_class(**param_config)) 174 | return Compose(component_list) if len(component_list) > 0 else None 175 | -------------------------------------------------------------------------------- /src/structure/sampler.py: -------------------------------------------------------------------------------- 1 | import bisect 2 | import copy 3 | from collections import defaultdict 4 | 5 | import numpy as np 6 | import torch 7 | import torch.utils.data 8 | import torchvision 9 | from PIL import Image 10 | from torch.utils.data.sampler import BatchSampler, Sampler 11 | from torch.utils.model_zoo import tqdm 12 | 13 | 14 | class GroupedBatchSampler(BatchSampler): 15 | """ 16 | Wraps another sampler to yield a mini-batch of indices. 17 | It enforces that the batch only contain elements from the same group. 18 | It also tries to provide mini-batches which follows an ordering which is 19 | as close as possible to the ordering from the original sampler. 20 | Arguments: 21 | sampler (Sampler): Base sampler. 22 | group_ids (list[int]): If the sampler produces indices in range [0, N), 23 | `group_ids` must be a list of `N` ints which contains the group id of each sample. 24 | The group ids must be a continuous set of integers starting from 25 | 0, i.e. they must be in the range [0, num_groups). 26 | batch_size (int): Size of mini-batch. 27 | """ 28 | def __init__(self, sampler, group_ids, batch_size): 29 | if not isinstance(sampler, Sampler): 30 | raise ValueError( 31 | 'sampler should be an instance of ' 32 | 'torch.utils.data.Sampler, but got sampler={}'.format(sampler) 33 | ) 34 | self.sampler = sampler 35 | self.group_ids = group_ids 36 | self.batch_size = batch_size 37 | 38 | def __iter__(self): 39 | buffer_per_group = defaultdict(list) 40 | samples_per_group = defaultdict(list) 41 | 42 | num_batches = 0 43 | for idx in self.sampler: 44 | group_id = self.group_ids[idx] 45 | buffer_per_group[group_id].append(idx) 46 | samples_per_group[group_id].append(idx) 47 | if len(buffer_per_group[group_id]) == self.batch_size: 48 | yield buffer_per_group[group_id] 49 | num_batches += 1 50 | del buffer_per_group[group_id] 51 | assert len(buffer_per_group[group_id]) < self.batch_size 52 | 53 | # now we have run out of elements that satisfy 54 | # the group criteria, let's return the remaining 55 | # elements so that the size of the sampler is 56 | # deterministic 57 | expected_num_batches = len(self) 58 | num_remaining = expected_num_batches - num_batches 59 | if num_remaining > 0: 60 | # for the remaining batches, take first the buffers with largest number 61 | # of elements 62 | for group_id, _ in sorted(buffer_per_group.items(), 63 | key=lambda x: len(x[1]), reverse=True): 64 | remaining = self.batch_size - len(buffer_per_group[group_id]) 65 | buffer_per_group[group_id].extend( 66 | samples_per_group[group_id][:remaining]) 67 | assert len(buffer_per_group[group_id]) == self.batch_size 68 | yield buffer_per_group[group_id] 69 | num_remaining -= 1 70 | if num_remaining == 0: 71 | break 72 | assert num_remaining == 0 73 | 74 | def __len__(self): 75 | return len(self.sampler) // self.batch_size 76 | 77 | 78 | class _SubsetSampler(Sampler): 79 | def __init__(self, indices): 80 | self.indices = indices 81 | 82 | def __iter__(self): 83 | return iter(self.indices) 84 | 85 | def __len__(self): 86 | return len(self.indices) 87 | 88 | 89 | def _compute_aspect_ratios_slow(dataset, indices=None): 90 | print('Your dataset doesn\'t support the fast path for ' 91 | 'computing the aspect ratios, so will iterate over ' 92 | 'the full dataset and load every image instead. ' 93 | 'This might take some time...') 94 | if indices is None: 95 | indices = range(len(dataset)) 96 | 97 | sampler = _SubsetSampler(indices) 98 | data_loader = torch.utils.data.DataLoader( 99 | dataset, batch_size=1, sampler=sampler, 100 | num_workers=14, # you might want to increase it for faster processing 101 | collate_fn=lambda x: x[0]) 102 | aspect_ratios = [] 103 | with tqdm(total=len(dataset)) as pbar: 104 | for _i, (img, _) in enumerate(data_loader): 105 | pbar.update(1) 106 | height, width = img.shape[-2:] 107 | aspect_ratio = float(width) / float(height) 108 | aspect_ratios.append(aspect_ratio) 109 | return aspect_ratios 110 | 111 | 112 | def _compute_aspect_ratios_custom_dataset(dataset, indices=None): 113 | if indices is None: 114 | indices = range(len(dataset)) 115 | aspect_ratios = [] 116 | for i in indices: 117 | height, width = dataset.get_height_and_width(i) 118 | aspect_ratio = float(width) / float(height) 119 | aspect_ratios.append(aspect_ratio) 120 | return aspect_ratios 121 | 122 | 123 | def _compute_aspect_ratios_coco_dataset(dataset, indices=None): 124 | if indices is None: 125 | indices = range(len(dataset)) 126 | aspect_ratios = [] 127 | for i in indices: 128 | img_info = dataset.coco.imgs[dataset.ids[i]] 129 | aspect_ratio = float(img_info['width']) / float(img_info['height']) 130 | aspect_ratios.append(aspect_ratio) 131 | return aspect_ratios 132 | 133 | 134 | def _compute_aspect_ratios_voc_dataset(dataset, indices=None): 135 | if indices is None: 136 | indices = range(len(dataset)) 137 | aspect_ratios = [] 138 | for i in indices: 139 | # this doesn't load the data into memory, because PIL loads it lazily 140 | width, height = Image.open(dataset.images[i]).size 141 | aspect_ratio = float(width) / float(height) 142 | aspect_ratios.append(aspect_ratio) 143 | return aspect_ratios 144 | 145 | 146 | def _compute_aspect_ratios_subset_dataset(dataset, indices=None): 147 | if indices is None: 148 | indices = range(len(dataset)) 149 | 150 | ds_indices = [dataset.indices[i] for i in indices] 151 | return compute_aspect_ratios(dataset.dataset, ds_indices) 152 | 153 | 154 | def compute_aspect_ratios(dataset, indices=None): 155 | if hasattr(dataset, 'get_height_and_width'): 156 | return _compute_aspect_ratios_custom_dataset(dataset, indices) 157 | 158 | if isinstance(dataset, torchvision.datasets.CocoDetection): 159 | return _compute_aspect_ratios_coco_dataset(dataset, indices) 160 | 161 | if isinstance(dataset, torchvision.datasets.VOCDetection): 162 | return _compute_aspect_ratios_voc_dataset(dataset, indices) 163 | 164 | if isinstance(dataset, torch.utils.data.Subset): 165 | return _compute_aspect_ratios_subset_dataset(dataset, indices) 166 | 167 | # slow path 168 | return _compute_aspect_ratios_slow(dataset, indices) 169 | 170 | 171 | def _quantize(x, bins): 172 | bins = copy.deepcopy(bins) 173 | bins = sorted(bins) 174 | quantized = list(map(lambda y: bisect.bisect_right(bins, y), x)) 175 | return quantized 176 | 177 | 178 | def create_aspect_ratio_groups(dataset, k=0): 179 | aspect_ratios = compute_aspect_ratios(dataset) 180 | bins = (2 ** np.linspace(-1, 1, 2 * k + 1)).tolist() if k > 0 else [1.0] 181 | groups = _quantize(aspect_ratios, bins) 182 | # count number of elements per group 183 | counts = np.unique(groups, return_counts=True)[1] 184 | fbins = [0] + bins + [np.inf] 185 | print('Using {} as bins for aspect ratio quantization'.format(fbins)) 186 | print('Count of instances per bin: {}'.format(counts)) 187 | return groups 188 | -------------------------------------------------------------------------------- /src/mimic_runner.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import datetime 3 | import time 4 | 5 | import torch 6 | from torch import distributed as dist 7 | from torch.nn import DataParallel 8 | from torch.nn.parallel.distributed import DistributedDataParallel 9 | 10 | from distillation.tool import DistillationBox 11 | from models import load_ckpt, get_model, save_ckpt 12 | from myutils.common import file_util, yaml_util 13 | from myutils.pytorch import func_util, module_util 14 | from utils import data_util, main_util, misc_util 15 | 16 | 17 | def get_argparser(): 18 | argparser = argparse.ArgumentParser(description='Mimic Runner') 19 | argparser.add_argument('--config', required=True, help='yaml file path') 20 | argparser.add_argument('--device', default='cuda', help='device') 21 | argparser.add_argument('--json', help='dictionary to overwrite config') 22 | argparser.add_argument('-distill', action='store_true', help='distill a teacher model') 23 | argparser.add_argument('-skip_teacher_eval', action='store_true', help='skip teacher model evaluation in testing') 24 | argparser.add_argument('-transform_bottleneck', action='store_true', 25 | help='use bottleneck transformer (if defined in yaml) in testing') 26 | # distributed training parameters 27 | argparser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') 28 | argparser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') 29 | return argparser 30 | 31 | 32 | def freeze_modules(student_model, student_model_config): 33 | for student_path in student_model_config['frozen_modules']: 34 | student_module = module_util.get_module(student_model, student_path) 35 | module_util.freeze_module_params(student_module) 36 | 37 | 38 | def distill_model(distillation_box, data_loader, optimizer, log_freq, device, epoch): 39 | metric_logger = misc_util.MetricLogger(delimiter=' ') 40 | metric_logger.add_meter('lr', misc_util.SmoothedValue(window_size=1, fmt='{value:.6f}')) 41 | header = 'Epoch: [{}]'.format(epoch) 42 | lr_scheduler = None 43 | if epoch == 0: 44 | warmup_factor = 1.0 / 1000.0 45 | warmup_iters = min(1000, len(data_loader) - 1) 46 | lr_scheduler = main_util.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor) 47 | 48 | for images, targets in metric_logger.log_every(data_loader, log_freq, header): 49 | images = list(image.to(device) for image in images) 50 | targets = [{k: v.to(device) for k, v in t.items()} for t in targets] 51 | loss = distillation_box(images, targets) 52 | optimizer.zero_grad() 53 | loss.backward() 54 | optimizer.step() 55 | if lr_scheduler is not None: 56 | lr_scheduler.step() 57 | 58 | metric_logger.update(loss=loss) 59 | metric_logger.update(lr=optimizer.param_groups[0]['lr']) 60 | 61 | 62 | def distill(teacher_model, student_model, train_sampler, train_data_loader, val_data_loader, 63 | device, distributed, distill_backbone_only, config, args): 64 | train_config = config['train'] 65 | distillation_box = DistillationBox(teacher_model, student_model, train_config['criterion']) 66 | ckpt_file_path = config['student_model']['ckpt'] 67 | optim_config = train_config['optimizer'] 68 | optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) 69 | scheduler_config = train_config['scheduler'] 70 | lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) 71 | use_bottleneck_transformer = args.transform_bottleneck 72 | best_val_map = 0.0 73 | if file_util.check_if_exists(ckpt_file_path): 74 | best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) 75 | 76 | num_epochs = train_config['num_epochs'] 77 | log_freq = train_config['log_freq'] 78 | teacher_model_without_dp = teacher_model.module if isinstance(teacher_model, DataParallel) else teacher_model 79 | student_model_without_ddp =\ 80 | student_model.module if isinstance(student_model, DistributedDataParallel) else student_model 81 | start_time = time.time() 82 | for epoch in range(num_epochs): 83 | if distributed: 84 | train_sampler.set_epoch(epoch) 85 | 86 | teacher_model.eval() 87 | student_model.train() 88 | teacher_model_without_dp.distill_backbone_only = distill_backbone_only 89 | student_model_without_ddp.distill_backbone_only = distill_backbone_only 90 | student_model_without_ddp.backbone.body.layer1.use_bottleneck_transformer = False 91 | distill_model(distillation_box, train_data_loader, optimizer, log_freq, device, epoch) 92 | student_model_without_ddp.distill_backbone_only = False 93 | student_model_without_ddp.backbone.body.layer1.use_bottleneck_transformer = use_bottleneck_transformer 94 | coco_evaluator = main_util.evaluate(student_model, val_data_loader, device=device) 95 | # Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] 96 | val_map = coco_evaluator.coco_eval['bbox'].stats[0] 97 | if val_map > best_val_map and misc_util.is_main_process(): 98 | print('Updating ckpt (Best BBox mAP: {:.4f} -> {:.4f})'.format(best_val_map, val_map)) 99 | best_val_map = val_map 100 | save_ckpt(student_model_without_ddp, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path) 101 | lr_scheduler.step() 102 | 103 | dist.barrier() 104 | total_time = time.time() - start_time 105 | total_time_str = str(datetime.timedelta(seconds=int(total_time))) 106 | print('Training time {}'.format(total_time_str)) 107 | 108 | 109 | def evaluate(teacher_model, student_model, test_data_loader, device, student_only, use_bottleneck_transformer): 110 | teacher_model_without_dp = teacher_model.module if isinstance(teacher_model, DataParallel) else teacher_model 111 | student_model_without_ddp =\ 112 | student_model.module if isinstance(student_model, DistributedDataParallel) else student_model 113 | teacher_model_without_dp.distill_backbone_only = False 114 | student_model_without_ddp.distill_backbone_only = False 115 | student_model_without_ddp.backbone.body.layer1.use_bottleneck_transformer = use_bottleneck_transformer 116 | if not student_only: 117 | print('[Teacher model]') 118 | main_util.evaluate(teacher_model, test_data_loader, device=device) 119 | 120 | print('\n[Student model]') 121 | main_util.evaluate(student_model, test_data_loader, device=device) 122 | 123 | 124 | def main(args): 125 | config = yaml_util.load_yaml_file(args.config) 126 | if args.json is not None: 127 | main_util.overwrite_config(config, args.json) 128 | 129 | distributed, device_ids = main_util.init_distributed_mode(args.world_size, args.dist_url) 130 | device = torch.device(args.device if torch.cuda.is_available() else 'cpu') 131 | teacher_model = get_model(config['teacher_model'], device) 132 | module_util.freeze_module_params(teacher_model) 133 | student_model_config = config['student_model'] 134 | student_model = get_model(student_model_config, device) 135 | freeze_modules(student_model, student_model_config) 136 | print('Updatable parameters: {}'.format(module_util.get_updatable_param_names(student_model))) 137 | distill_backbone_only = student_model_config['distill_backbone_only'] 138 | train_config = config['train'] 139 | train_sampler, train_data_loader, val_data_loader, test_data_loader = \ 140 | data_util.get_coco_data_loaders(config['dataset'], train_config['batch_size'], distributed) 141 | if distributed: 142 | teacher_model = DataParallel(teacher_model, device_ids=device_ids) 143 | student_model = DistributedDataParallel(student_model, device_ids=device_ids) 144 | 145 | if args.distill: 146 | distill(teacher_model, student_model, train_sampler, train_data_loader, val_data_loader, 147 | device, distributed, distill_backbone_only, config, args) 148 | load_ckpt(config['student_model']['ckpt'], 149 | model=student_model.module if isinstance(student_model, DistributedDataParallel) else student_model) 150 | evaluate(teacher_model, student_model, test_data_loader, device, 151 | args.skip_teacher_eval, args.transform_bottleneck) 152 | 153 | 154 | if __name__ == '__main__': 155 | parser = get_argparser() 156 | main(parser.parse_args()) 157 | -------------------------------------------------------------------------------- /src/utils/misc_util.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import pickle 3 | import time 4 | from collections import defaultdict, deque 5 | 6 | import torch 7 | import torch.distributed as dist 8 | 9 | 10 | class SmoothedValue(object): 11 | """Track a series of values and provide access to smoothed values over a 12 | window or the global series average. 13 | """ 14 | 15 | def __init__(self, window_size=20, fmt=None): 16 | if fmt is None: 17 | fmt = '{median:.4f} ({global_avg:.4f})' 18 | self.deque = deque(maxlen=window_size) 19 | self.total = 0.0 20 | self.count = 0 21 | self.fmt = fmt 22 | 23 | def update(self, value, n=1): 24 | self.deque.append(value) 25 | self.count += n 26 | self.total += value * n 27 | 28 | def synchronize_between_processes(self): 29 | """ 30 | Warning: does not synchronize the deque! 31 | """ 32 | if not is_dist_avail_and_initialized(): 33 | return 34 | t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda') 35 | dist.barrier() 36 | dist.all_reduce(t) 37 | t = t.tolist() 38 | self.count = int(t[0]) 39 | self.total = t[1] 40 | 41 | @property 42 | def median(self): 43 | d = torch.tensor(list(self.deque)) 44 | return d.median().item() 45 | 46 | @property 47 | def avg(self): 48 | d = torch.tensor(list(self.deque), dtype=torch.float32) 49 | return d.mean().item() 50 | 51 | @property 52 | def global_avg(self): 53 | return self.total / self.count 54 | 55 | @property 56 | def max(self): 57 | return max(self.deque) 58 | 59 | @property 60 | def value(self): 61 | return self.deque[-1] 62 | 63 | def __str__(self): 64 | return self.fmt.format( 65 | median=self.median, 66 | avg=self.avg, 67 | global_avg=self.global_avg, 68 | max=self.max, 69 | value=self.value) 70 | 71 | 72 | def all_gather(data): 73 | """ 74 | Run all_gather on arbitrary picklable data (not necessarily tensors) 75 | Args: 76 | data: any picklable object 77 | Returns: 78 | list[data]: list of data gathered from each rank 79 | """ 80 | world_size = get_world_size() 81 | if world_size == 1: 82 | return [data] 83 | 84 | # serialized to a Tensor 85 | buffer = pickle.dumps(data) 86 | storage = torch.ByteStorage.from_buffer(buffer) 87 | tensor = torch.ByteTensor(storage).to('cuda') 88 | 89 | # obtain Tensor size of each rank 90 | local_size = torch.tensor([tensor.numel()], device='cuda') 91 | size_list = [torch.tensor([0], device='cuda') for _ in range(world_size)] 92 | dist.all_gather(size_list, local_size) 93 | size_list = [int(size.item()) for size in size_list] 94 | max_size = max(size_list) 95 | 96 | # receiving Tensor from all ranks 97 | # we pad the tensor because torch all_gather does not support 98 | # gathering tensors of different shapes 99 | tensor_list = [] 100 | for _ in size_list: 101 | tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device='cuda')) 102 | if local_size != max_size: 103 | padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device='cuda') 104 | tensor = torch.cat((tensor, padding), dim=0) 105 | dist.all_gather(tensor_list, tensor) 106 | 107 | data_list = [] 108 | for size, tensor in zip(size_list, tensor_list): 109 | buffer = tensor.cpu().numpy().tobytes()[:size] 110 | data_list.append(pickle.loads(buffer)) 111 | 112 | return data_list 113 | 114 | 115 | def reduce_dict(input_dict, average=True): 116 | """ 117 | Args: 118 | input_dict (dict): all the values will be reduced 119 | average (bool): whether to do average or sum 120 | Reduce the values in the dictionary from all processes so that all processes 121 | have the averaged results. Returns a dict with the same fields as 122 | input_dict, after reduction. 123 | """ 124 | world_size = get_world_size() 125 | if world_size < 2: 126 | return input_dict 127 | with torch.no_grad(): 128 | names = [] 129 | values = [] 130 | # sort the keys so that they are consistent across processes 131 | for k in sorted(input_dict.keys()): 132 | names.append(k) 133 | values.append(input_dict[k]) 134 | values = torch.stack(values, dim=0) 135 | dist.all_reduce(values) 136 | if average: 137 | values /= world_size 138 | reduced_dict = {k: v for k, v in zip(names, values)} 139 | return reduced_dict 140 | 141 | 142 | class MetricLogger(object): 143 | def __init__(self, delimiter='\t'): 144 | self.meters = defaultdict(SmoothedValue) 145 | self.delimiter = delimiter 146 | 147 | def update(self, **kwargs): 148 | for k, v in kwargs.items(): 149 | if isinstance(v, torch.Tensor): 150 | v = v.item() 151 | assert isinstance(v, (float, int)) 152 | self.meters[k].update(v) 153 | 154 | def __getattr__(self, attr): 155 | if attr in self.meters: 156 | return self.meters[attr] 157 | if attr in self.__dict__: 158 | return self.__dict__[attr] 159 | raise AttributeError('`{}` object has no attribute `{}`'.format( 160 | type(self).__name__, attr)) 161 | 162 | def __str__(self): 163 | loss_str = [] 164 | for name, meter in self.meters.items(): 165 | loss_str.append( 166 | '{}: {}'.format(name, str(meter)) 167 | ) 168 | return self.delimiter.join(loss_str) 169 | 170 | def synchronize_between_processes(self): 171 | for meter in self.meters.values(): 172 | meter.synchronize_between_processes() 173 | 174 | def add_meter(self, name, meter): 175 | self.meters[name] = meter 176 | 177 | def log_every(self, iterable, print_freq, header=None): 178 | i = 0 179 | if not header: 180 | header = '' 181 | start_time = time.time() 182 | end = time.time() 183 | iter_time = SmoothedValue(fmt='{avg:.4f}') 184 | data_time = SmoothedValue(fmt='{avg:.4f}') 185 | space_fmt = ':' + str(len(str(len(iterable)))) + 'd' 186 | if torch.cuda.is_available(): 187 | log_msg = self.delimiter.join([ 188 | header, 189 | '[{0' + space_fmt + '}/{1}]', 190 | 'eta: {eta}', 191 | '{meters}', 192 | 'time: {time}', 193 | 'data: {data}', 194 | 'max mem: {memory:.0f}' 195 | ]) 196 | else: 197 | log_msg = self.delimiter.join([ 198 | header, 199 | '[{0' + space_fmt + '}/{1}]', 200 | 'eta: {eta}', 201 | '{meters}', 202 | 'time: {time}', 203 | 'data: {data}' 204 | ]) 205 | MB = 1024.0 * 1024.0 206 | for obj in iterable: 207 | data_time.update(time.time() - end) 208 | yield obj 209 | iter_time.update(time.time() - end) 210 | if i % print_freq == 0 or i == len(iterable) - 1: 211 | eta_seconds = iter_time.global_avg * (len(iterable) - i) 212 | eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) 213 | if torch.cuda.is_available(): 214 | print(log_msg.format( 215 | i, len(iterable), eta=eta_string, 216 | meters=str(self), 217 | time=str(iter_time), data=str(data_time), 218 | memory=torch.cuda.max_memory_allocated() / MB)) 219 | else: 220 | print(log_msg.format( 221 | i, len(iterable), eta=eta_string, 222 | meters=str(self), 223 | time=str(iter_time), data=str(data_time))) 224 | i += 1 225 | end = time.time() 226 | total_time = time.time() - start_time 227 | total_time_str = str(datetime.timedelta(seconds=int(total_time))) 228 | print('{} Total time: {} ({:.4f} s / it)'.format( 229 | header, total_time_str, total_time / len(iterable))) 230 | 231 | 232 | def collate_fn(batch): 233 | return tuple(zip(*batch)) 234 | 235 | 236 | def is_dist_avail_and_initialized(): 237 | if not dist.is_available(): 238 | return False 239 | if not dist.is_initialized(): 240 | return False 241 | return True 242 | 243 | 244 | def get_world_size(): 245 | if not is_dist_avail_and_initialized(): 246 | return 1 247 | return dist.get_world_size() 248 | 249 | 250 | def get_rank(): 251 | if not is_dist_avail_and_initialized(): 252 | return 0 253 | return dist.get_rank() 254 | 255 | 256 | def is_main_process(): 257 | return get_rank() == 0 258 | 259 | 260 | def save_on_master(*args, **kwargs): 261 | if is_main_process(): 262 | torch.save(*args, **kwargs) 263 | -------------------------------------------------------------------------------- /src/ext_runner.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import datetime 3 | import math 4 | import sys 5 | import time 6 | 7 | import numpy as np 8 | import pandas as pd 9 | import torch 10 | from sklearn import metrics 11 | from torch import distributed as dist 12 | from torch import nn 13 | 14 | from models import get_model, load_ckpt, save_ckpt 15 | from models.ext.backbone import check_if_valid_target 16 | from myutils.common import file_util, yaml_util 17 | from myutils.pytorch import func_util, module_util 18 | from utils import data_util, main_util, misc_util 19 | 20 | 21 | def get_argparser(): 22 | argparser = argparse.ArgumentParser(description=__doc__) 23 | argparser.add_argument('--config', required=True, help='yaml config file') 24 | argparser.add_argument('--device', default='cuda', help='device') 25 | argparser.add_argument('--json', help='dictionary to overwrite config') 26 | argparser.add_argument('--min_recall', type=float, default=0.9, help='minimum recall to decide a threshold') 27 | argparser.add_argument('-train', action='store_true', help='train a model') 28 | # distributed training parameters 29 | argparser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') 30 | argparser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') 31 | return argparser 32 | 33 | 34 | def convert_target2ext_targets(targets, device): 35 | ext_targets = [1 if check_if_valid_target(target) else 0 for target in targets] 36 | return torch.LongTensor(ext_targets).to(device) 37 | 38 | 39 | def train_model(model, optimizer, data_loader, device, epoch, log_freq): 40 | model.train() 41 | metric_logger = misc_util.MetricLogger(delimiter=' ') 42 | metric_logger.add_meter('lr', misc_util.SmoothedValue(window_size=1, fmt='{value:.6f}')) 43 | header = 'Epoch: [{}]'.format(epoch) 44 | lr_scheduler = None 45 | if epoch == 0: 46 | warmup_factor = 1.0 / 1000.0 47 | warmup_iters = min(1000, len(data_loader) - 1) 48 | lr_scheduler = main_util.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor) 49 | 50 | for images, targets in metric_logger.log_every(data_loader, log_freq, header): 51 | images = list(image.to(device) for image in images) 52 | targets = [{k: v.to(device) for k, v in t.items()} for t in targets] 53 | ext_logits = model(images, targets) 54 | ext_targets = convert_target2ext_targets(targets, device) 55 | ext_cls_loss = nn.functional.cross_entropy(ext_logits, ext_targets) 56 | loss_dict = {'loss_ext_classifier': ext_cls_loss} 57 | losses = sum(loss for loss in loss_dict.values()) 58 | 59 | # reduce losses over all GPUs for logging purposes 60 | loss_dict_reduced = misc_util.reduce_dict(loss_dict) 61 | losses_reduced = sum(loss for loss in loss_dict_reduced.values()) 62 | loss_value = losses_reduced.item() 63 | 64 | if not math.isfinite(loss_value): 65 | print('Loss is {}, stopping training'.format(loss_value)) 66 | print(loss_dict_reduced) 67 | sys.exit(1) 68 | 69 | optimizer.zero_grad() 70 | losses.backward() 71 | optimizer.step() 72 | if lr_scheduler is not None: 73 | lr_scheduler.step() 74 | 75 | metric_logger.update(loss=losses_reduced, **loss_dict_reduced) 76 | metric_logger.update(lr=optimizer.param_groups[0]['lr']) 77 | 78 | 79 | def evaluate(model, data_loader, device, min_recall, split_name='Validation'): 80 | model.eval() 81 | correct_count = 0 82 | pos_correct_count = 0 83 | pos_count = 0 84 | prob_list = list() 85 | label_list = list() 86 | with torch.no_grad(): 87 | for images, targets in data_loader: 88 | images = list(image.to(device) for image in images) 89 | targets = [{k: v.to(device) for k, v in t.items()} for t in targets] 90 | ext_logits = model(images, targets) 91 | ext_targets = convert_target2ext_targets(targets, device) 92 | prob_list.append(ext_logits[:, 1].cpu().numpy()) 93 | label_list.append(ext_targets.cpu().numpy()) 94 | preds = ext_logits.argmax(dim=1) 95 | correct_count += preds.eq(ext_targets).sum().item() 96 | pos_correct_count += preds[ext_targets.nonzero().flatten()].sum().item() 97 | pos_count += ext_targets.sum().item() 98 | 99 | num_samples = len(data_loader.dataset) 100 | accuracy = correct_count / num_samples 101 | recall = pos_correct_count / pos_count 102 | specificity = (correct_count - pos_correct_count) / (num_samples - pos_count) 103 | probs = np.concatenate(prob_list) 104 | labels = np.concatenate(label_list) 105 | roc_auc = metrics.roc_auc_score(labels, probs) 106 | print('[{}]'.format(split_name)) 107 | print('\tAccuracy: {:.4f} ({} / {})'.format(accuracy, correct_count, num_samples)) 108 | print('\tRecall: {:.4f} ({} / {})'.format(recall, pos_correct_count, pos_count)) 109 | print('\tSpecificity: {:.4f} ({} / {})'.format(specificity, correct_count - pos_correct_count, 110 | num_samples - pos_count)) 111 | print('\tROC-AUC: {:.4f}'.format(roc_auc)) 112 | if split_name == 'Test': 113 | fprs, tprs, thrs = metrics.roc_curve(labels, probs, pos_label=1) 114 | idx = np.searchsorted(tprs, min_recall) 115 | 116 | data_frame =\ 117 | pd.DataFrame(np.array([thrs[idx:], tprs[idx:], fprs[idx:]]).T, columns=['Threshold', 'TPR (Recall)', 'FPR']) 118 | with pd.option_context('display.max_rows', None, 'display.max_columns', None): 119 | print(data_frame) 120 | return roc_auc 121 | 122 | 123 | def train(model, ext_classifier, train_sampler, train_data_loader, val_data_loader, device, distributed, 124 | config, args, ckpt_file_path): 125 | train_config = config['train'] 126 | optim_config = train_config['optimizer'] 127 | optimizer = func_util.get_optimizer(ext_classifier, optim_config['type'], optim_config['params']) 128 | scheduler_config = train_config['scheduler'] 129 | lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) 130 | best_val_roc_auc = 0.0 131 | if file_util.check_if_exists(ckpt_file_path): 132 | best_val_roc_auc, _, _ =\ 133 | load_ckpt(ckpt_file_path, model=ext_classifier, optimizer=optimizer, lr_scheduler=lr_scheduler) 134 | 135 | num_epochs = train_config['num_epochs'] 136 | log_freq = train_config['log_freq'] 137 | start_time = time.time() 138 | for epoch in range(num_epochs): 139 | if distributed: 140 | train_sampler.set_epoch(epoch) 141 | 142 | train_model(model, optimizer, train_data_loader, device, epoch, log_freq) 143 | lr_scheduler.step() 144 | 145 | # evaluate after every epoch 146 | val_roc_auc = evaluate(model, val_data_loader, device, min_recall=args.min_recall, split_name='Validation') 147 | if val_roc_auc > best_val_roc_auc: 148 | print('Updating ckpt (Best ROC-AUC: {:.4f} -> {:.4f})'.format(best_val_roc_auc, val_roc_auc)) 149 | best_val_roc_auc = val_roc_auc 150 | save_ckpt(ext_classifier, optimizer, lr_scheduler, best_val_roc_auc, config, args, ckpt_file_path) 151 | 152 | dist.barrier() 153 | total_time = time.time() - start_time 154 | total_time_str = str(datetime.timedelta(seconds=int(total_time))) 155 | print('Training time {}'.format(total_time_str)) 156 | 157 | 158 | def main(args): 159 | distributed, device_ids = main_util.init_distributed_mode(args.world_size, args.dist_url) 160 | config = yaml_util.load_yaml_file(args.config) 161 | if args.json is not None: 162 | main_util.overwrite_config(config, args.json) 163 | 164 | device = torch.device(args.device) 165 | print(args) 166 | print('Loading data') 167 | train_config = config['train'] 168 | train_sampler, train_data_loader, val_data_loader, test_data_loader =\ 169 | data_util.get_coco_data_loaders(config['dataset'], train_config['batch_size'], distributed) 170 | 171 | print('Creating model') 172 | model_config = config['model'] 173 | model = get_model(model_config, device, strict=False) 174 | module_util.freeze_module_params(model) 175 | ext_classifier = model.get_ext_classifier() 176 | module_util.unfreeze_module_params(ext_classifier) 177 | print('Updatable parameters: {}'.format(module_util.get_updatable_param_names(model))) 178 | model.train_ext() 179 | if distributed: 180 | model = nn.parallel.DistributedDataParallel(model, device_ids=device_ids) 181 | 182 | if args.train: 183 | print('Start training') 184 | start_time = time.time() 185 | ckpt_file_path = model_config['backbone']['ext_config']['ckpt'] 186 | train(model, ext_classifier, train_sampler, train_data_loader, val_data_loader, device, distributed, 187 | config, args, ckpt_file_path) 188 | total_time = time.time() - start_time 189 | total_time_str = str(datetime.timedelta(seconds=int(total_time))) 190 | print('Training time {}'.format(total_time_str)) 191 | load_ckpt(ckpt_file_path, model=ext_classifier) 192 | evaluate(model, test_data_loader, device=device, min_recall=args.min_recall, split_name='Test') 193 | 194 | 195 | if __name__ == '__main__': 196 | parser = get_argparser() 197 | main(parser.parse_args()) 198 | -------------------------------------------------------------------------------- /src/utils/coco_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import os 3 | from io import BytesIO 4 | 5 | import torch 6 | import torch.utils.data 7 | import torchvision 8 | from PIL import Image 9 | from pycocotools import mask as coco_mask 10 | from pycocotools.coco import COCO 11 | 12 | from structure.transformer import Compose 13 | 14 | 15 | class FilterAndRemapCocoCategories(object): 16 | def __init__(self, categories, remap=True): 17 | self.categories = categories 18 | self.remap = remap 19 | 20 | def __call__(self, image, target): 21 | anno = target['annotations'] 22 | anno = [obj for obj in anno if obj['category_id'] in self.categories] 23 | if not self.remap: 24 | target['annotations'] = anno 25 | return image, target 26 | anno = copy.deepcopy(anno) 27 | for obj in anno: 28 | obj['category_id'] = self.categories.index(obj['category_id']) 29 | target['annotations'] = anno 30 | return image, target 31 | 32 | 33 | def convert_coco_poly_to_mask(segmentations, height, width): 34 | masks = [] 35 | for polygons in segmentations: 36 | rles = coco_mask.frPyObjects(polygons, height, width) 37 | mask = coco_mask.decode(rles) 38 | if len(mask.shape) < 3: 39 | mask = mask[..., None] 40 | mask = torch.as_tensor(mask, dtype=torch.uint8) 41 | mask = mask.any(dim=2) 42 | masks.append(mask) 43 | if masks: 44 | masks = torch.stack(masks, dim=0) 45 | else: 46 | masks = torch.zeros((0, height, width), dtype=torch.uint8) 47 | return masks 48 | 49 | 50 | class ConvertCocoPolysToMask(object): 51 | def __call__(self, image, target): 52 | w, h = image.size 53 | 54 | image_id = target['image_id'] 55 | image_id = torch.tensor([image_id]) 56 | 57 | anno = target['annotations'] 58 | 59 | anno = [obj for obj in anno if obj['iscrowd'] == 0] 60 | 61 | boxes = [obj['bbox'] for obj in anno] 62 | # guard against no boxes via resizing 63 | boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4) 64 | boxes[:, 2:] += boxes[:, :2] 65 | boxes[:, 0::2].clamp_(min=0, max=w) 66 | boxes[:, 1::2].clamp_(min=0, max=h) 67 | 68 | classes = [obj['category_id'] for obj in anno] 69 | classes = torch.tensor(classes, dtype=torch.int64) 70 | 71 | segmentations = [obj['segmentation'] for obj in anno] 72 | masks = convert_coco_poly_to_mask(segmentations, h, w) 73 | 74 | keypoints = None 75 | if anno and 'keypoints' in anno[0]: 76 | keypoints = [obj['keypoints'] for obj in anno] 77 | keypoints = torch.as_tensor(keypoints, dtype=torch.float32) 78 | num_keypoints = keypoints.shape[0] 79 | if num_keypoints: 80 | keypoints = keypoints.view(num_keypoints, -1, 3) 81 | 82 | keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0]) 83 | boxes = boxes[keep] 84 | classes = classes[keep] 85 | masks = masks[keep] 86 | if keypoints is not None: 87 | keypoints = keypoints[keep] 88 | 89 | target = {} 90 | target['boxes'] = boxes 91 | target['labels'] = classes 92 | target['masks'] = masks 93 | target['image_id'] = image_id 94 | if keypoints is not None: 95 | target['keypoints'] = keypoints 96 | 97 | # for conversion to coco api 98 | area = torch.tensor([obj['area'] for obj in anno]) 99 | iscrowd = torch.tensor([obj['iscrowd'] for obj in anno]) 100 | target['area'] = area 101 | target['iscrowd'] = iscrowd 102 | 103 | return image, target 104 | 105 | 106 | def has_only_empty_bbox(anno): 107 | return all(any(o <= 1 for o in obj['bbox'][2:]) for obj in anno) 108 | 109 | 110 | def count_visible_keypoints(anno): 111 | return sum(sum(1 for v in ann['keypoints'][2::3] if v > 0) for ann in anno) 112 | 113 | 114 | def has_valid_annotation(anno, min_keypoints_per_image=10): 115 | # if it's empty, there is no annotation 116 | if len(anno) == 0: 117 | return False 118 | # if all boxes have close to zero area, there is no annotation 119 | if has_only_empty_bbox(anno): 120 | return False 121 | # keypoints task have a slight different criteria for considering 122 | # if an annotation is valid 123 | if 'keypoints' not in anno[0]: 124 | return True 125 | # for keypoint detection tasks, only consider valid images those 126 | # containing at least min_keypoints_per_image 127 | if count_visible_keypoints(anno) >= min_keypoints_per_image: 128 | return True 129 | return False 130 | 131 | 132 | def remove_images_without_annotations(dataset, cat_list=None): 133 | assert isinstance(dataset, torchvision.datasets.CocoDetection) 134 | ids = [] 135 | for ds_idx, img_id in enumerate(dataset.ids): 136 | ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None) 137 | anno = dataset.coco.loadAnns(ann_ids) 138 | if cat_list: 139 | anno = [obj for obj in anno if obj['category_id'] in cat_list] 140 | if has_valid_annotation(anno): 141 | ids.append(ds_idx) 142 | 143 | dataset = torch.utils.data.Subset(dataset, ids) 144 | return dataset 145 | 146 | 147 | def convert_to_coco_api(ds): 148 | coco_ds = COCO() 149 | # annotation IDs need to start at 1, not 0, see torchvision issue #1530 150 | ann_id = 1 151 | dataset = {'images': [], 'categories': [], 'annotations': []} 152 | categories = set() 153 | for img_idx in range(len(ds)): 154 | # find better way to get target 155 | # targets = ds.get_annotations(img_idx) 156 | img, targets = ds[img_idx] 157 | image_id = targets['image_id'].item() 158 | img_dict = {} 159 | img_dict['id'] = image_id 160 | img_dict['height'] = img.shape[-2] 161 | img_dict['width'] = img.shape[-1] 162 | dataset['images'].append(img_dict) 163 | bboxes = targets['boxes'] 164 | bboxes[:, 2:] -= bboxes[:, :2] 165 | bboxes = bboxes.tolist() 166 | labels = targets['labels'].tolist() 167 | areas = targets['area'].tolist() 168 | iscrowd = targets['iscrowd'].tolist() 169 | if 'masks' in targets: 170 | masks = targets['masks'] 171 | # make masks Fortran contiguous for coco_mask 172 | masks = masks.permute(0, 2, 1).contiguous().permute(0, 2, 1) 173 | if 'keypoints' in targets: 174 | keypoints = targets['keypoints'] 175 | keypoints = keypoints.reshape(keypoints.shape[0], -1).tolist() 176 | num_objs = len(bboxes) 177 | for i in range(num_objs): 178 | ann = {} 179 | ann['image_id'] = image_id 180 | ann['bbox'] = bboxes[i] 181 | ann['category_id'] = labels[i] 182 | categories.add(labels[i]) 183 | ann['area'] = areas[i] 184 | ann['iscrowd'] = iscrowd[i] 185 | ann['id'] = ann_id 186 | if 'masks' in targets: 187 | ann['segmentation'] = coco_mask.encode(masks[i].numpy()) 188 | if 'keypoints' in targets: 189 | ann['keypoints'] = keypoints[i] 190 | ann['num_keypoints'] = sum(k != 0 for k in keypoints[i][2::3]) 191 | dataset['annotations'].append(ann) 192 | ann_id += 1 193 | dataset['categories'] = [{'id': i} for i in sorted(categories)] 194 | coco_ds.dataset = dataset 195 | coco_ds.createIndex() 196 | return coco_ds 197 | 198 | 199 | def get_coco_api_from_dataset(dataset): 200 | for _ in range(10): 201 | if isinstance(dataset, torchvision.datasets.CocoDetection): 202 | break 203 | if isinstance(dataset, torch.utils.data.Subset): 204 | dataset = dataset.dataset 205 | if isinstance(dataset, torchvision.datasets.CocoDetection): 206 | return dataset.coco 207 | return convert_to_coco_api(dataset) 208 | 209 | 210 | class ExtCocoDetection(torchvision.datasets.CocoDetection): 211 | def __init__(self, img_folder, ann_file, transforms, jpeg_quality): 212 | super().__init__(img_folder, ann_file) 213 | self.additional_transforms = transforms 214 | self.jpeg_quality = jpeg_quality if jpeg_quality is not None and 1 <= jpeg_quality <= 95 else None 215 | 216 | def __getitem__(self, index): 217 | coco = self.coco 218 | img_id = self.ids[index] 219 | ann_ids = coco.getAnnIds(imgIds=img_id) 220 | target = coco.loadAnns(ann_ids) 221 | path = coco.loadImgs(img_id)[0]['file_name'] 222 | img = Image.open(os.path.join(self.root, path)).convert('RGB') 223 | if self.jpeg_quality is not None: 224 | img_buffer = BytesIO() 225 | img.save(img_buffer, 'JPEG', quality=self.jpeg_quality) 226 | img = Image.open(img_buffer) 227 | 228 | if self.transforms is not None: 229 | img, target = self.transforms(img, target) 230 | 231 | image_id = self.ids[index] 232 | target = dict(image_id=image_id, annotations=target) 233 | if self.additional_transforms is not None: 234 | img, target = self.additional_transforms(img, target) 235 | return img, target 236 | 237 | 238 | def get_coco(img_dir_path, ann_file_path, transforms, remove_non_annotated_imgs, jpeg_quality=None): 239 | t = [ConvertCocoPolysToMask()] 240 | if transforms is not None: 241 | t.append(transforms) 242 | 243 | transforms = Compose(t) 244 | dataset = ExtCocoDetection(os.path.expanduser(img_dir_path), os.path.expanduser(ann_file_path), 245 | transforms=transforms, jpeg_quality=jpeg_quality) 246 | if remove_non_annotated_imgs: 247 | dataset = remove_images_without_annotations(dataset) 248 | return dataset 249 | -------------------------------------------------------------------------------- /src/models/custom/resnet.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torchvision.models.resnet import BasicBlock, Bottleneck, conv1x1, model_urls 3 | from torchvision.models.utils import load_state_dict_from_url 4 | 5 | 6 | class CustomResNet(nn.Module): 7 | def __init__(self, block, layers, zero_init_residual=False, 8 | groups=1, width_per_group=64, replace_stride_with_dilation=None, 9 | norm_layer=None, layer1=None, layer2=None, layer3=None, layer4=None): 10 | super().__init__() 11 | if norm_layer is None: 12 | norm_layer = nn.BatchNorm2d 13 | self._norm_layer = norm_layer 14 | 15 | self.inplanes = 64 16 | self.dilation = 1 17 | if replace_stride_with_dilation is None: 18 | # each element in the tuple indicates if we should replace 19 | # the 2x2 stride with a dilated convolution instead 20 | replace_stride_with_dilation = [False, False, False] 21 | if len(replace_stride_with_dilation) != 3: 22 | raise ValueError("replace_stride_with_dilation should be None " 23 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) 24 | self.groups = groups 25 | self.base_width = width_per_group 26 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, 27 | bias=False) 28 | self.bn1 = norm_layer(self.inplanes) 29 | self.relu = nn.ReLU(inplace=True) 30 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 31 | if layer1 is None: 32 | self.layer1 = self._make_layer(block, 64, layers[0]) 33 | else: 34 | self.layer1 = layer1 35 | self.inplanes = 64 * block.expansion 36 | 37 | if layer2 is None: 38 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) 39 | else: 40 | self.layer2 = layer2 41 | self.inplanes = 128 * block.expansion 42 | 43 | if layer3 is None: 44 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) 45 | else: 46 | self.layer3 = layer3 47 | self.inplanes = 256 * block.expansion 48 | 49 | if layer4 is None: 50 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) 51 | else: 52 | self.layer4 = layer4 53 | self.inplanes = 512 * block.expansion 54 | 55 | for m in self.modules(): 56 | if isinstance(m, nn.Conv2d): 57 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 58 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): 59 | nn.init.constant_(m.weight, 1) 60 | nn.init.constant_(m.bias, 0) 61 | 62 | # Zero-initialize the last BN in each residual branch, 63 | # so that the residual branch starts with zeros, and each residual block behaves like an identity. 64 | # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 65 | if zero_init_residual: 66 | for m in self.modules(): 67 | if isinstance(m, Bottleneck): 68 | nn.init.constant_(m.bn3.weight, 0) 69 | elif isinstance(m, BasicBlock): 70 | nn.init.constant_(m.bn2.weight, 0) 71 | 72 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False): 73 | norm_layer = self._norm_layer 74 | downsample = None 75 | previous_dilation = self.dilation 76 | if dilate: 77 | self.dilation *= stride 78 | stride = 1 79 | if stride != 1 or self.inplanes != planes * block.expansion: 80 | downsample = nn.Sequential( 81 | conv1x1(self.inplanes, planes * block.expansion, stride), 82 | norm_layer(planes * block.expansion), 83 | ) 84 | 85 | layers = list() 86 | layers.append(block(self.inplanes, planes, stride, downsample, self.groups, 87 | self.base_width, previous_dilation, norm_layer)) 88 | self.inplanes = planes * block.expansion 89 | for _ in range(1, blocks): 90 | layers.append(block(self.inplanes, planes, groups=self.groups, 91 | base_width=self.base_width, dilation=self.dilation, 92 | norm_layer=norm_layer)) 93 | return nn.Sequential(*layers) 94 | 95 | def forward(self, x): 96 | x = self.conv1(x) 97 | x = self.bn1(x) 98 | x = self.relu(x) 99 | x = self.maxpool(x) 100 | 101 | x = self.layer1(x) 102 | x = self.layer2(x) 103 | x = self.layer3(x) 104 | x = self.layer4(x) 105 | return x 106 | 107 | 108 | def _custom_resnet(arch, block, layers, pretrained, progress, strict=False, **kwargs): 109 | model = CustomResNet(block, layers, **kwargs) 110 | if pretrained: 111 | print('Loading pretrained state dict of {}'.format(arch)) 112 | state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) 113 | model.load_state_dict(state_dict, strict=strict) 114 | return model 115 | 116 | 117 | def custom_resnet18(pretrained=False, progress=True, **kwargs): 118 | r"""ResNet-18 model from 119 | `"Deep Residual Learning for Image Recognition" `_ 120 | Args: 121 | pretrained (bool): If True, returns a model pre-trained on ImageNet 122 | progress (bool): If True, displays a progress bar of the download to stderr 123 | """ 124 | return _custom_resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs) 125 | 126 | 127 | def custom_resnet34(pretrained=False, progress=True, **kwargs): 128 | r"""ResNet-34 model from 129 | `"Deep Residual Learning for Image Recognition" `_ 130 | Args: 131 | pretrained (bool): If True, returns a model pre-trained on ImageNet 132 | progress (bool): If True, displays a progress bar of the download to stderr 133 | """ 134 | return _custom_resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs) 135 | 136 | 137 | def custom_resnet50(pretrained=False, progress=True, **kwargs): 138 | r"""ResNet-50 model from 139 | `"Deep Residual Learning for Image Recognition" `_ 140 | Args: 141 | pretrained (bool): If True, returns a model pre-trained on ImageNet 142 | progress (bool): If True, displays a progress bar of the download to stderr 143 | """ 144 | return _custom_resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) 145 | 146 | 147 | def custom_resnet101(pretrained=False, progress=True, **kwargs): 148 | r"""ResNet-101 model from 149 | `"Deep Residual Learning for Image Recognition" `_ 150 | Args: 151 | pretrained (bool): If True, returns a model pre-trained on ImageNet 152 | progress (bool): If True, displays a progress bar of the download to stderr 153 | """ 154 | return _custom_resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) 155 | 156 | 157 | def custom_resnet152(pretrained=False, progress=True, **kwargs): 158 | r"""ResNet-152 model from 159 | `"Deep Residual Learning for Image Recognition" `_ 160 | Args: 161 | pretrained (bool): If True, returns a model pre-trained on ImageNet 162 | progress (bool): If True, displays a progress bar of the download to stderr 163 | """ 164 | return _custom_resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs) 165 | 166 | 167 | def custom_resnext50_32x4d(pretrained=False, progress=True, **kwargs): 168 | r"""ResNeXt-50 32x4d model from 169 | `"Aggregated Residual Transformation for Deep Neural Networks" `_ 170 | Args: 171 | pretrained (bool): If True, returns a model pre-trained on ImageNet 172 | progress (bool): If True, displays a progress bar of the download to stderr 173 | """ 174 | kwargs['groups'] = 32 175 | kwargs['width_per_group'] = 4 176 | return _custom_resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) 177 | 178 | 179 | def cusom_resnext101_32x8d(pretrained=False, progress=True, **kwargs): 180 | r"""ResNeXt-101 32x8d model from 181 | `"Aggregated Residual Transformation for Deep Neural Networks" `_ 182 | Args: 183 | pretrained (bool): If True, returns a model pre-trained on ImageNet 184 | progress (bool): If True, displays a progress bar of the download to stderr 185 | """ 186 | kwargs['groups'] = 32 187 | kwargs['width_per_group'] = 8 188 | return _custom_resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) 189 | 190 | 191 | def custom_wide_resnet50_2(pretrained=False, progress=True, **kwargs): 192 | r"""Wide ResNet-50-2 model from 193 | `"Wide Residual Networks" `_ 194 | The model is the same as ResNet except for the bottleneck number of channels 195 | which is twice larger in every block. The number of channels in outer 1x1 196 | convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 197 | channels, and in Wide ResNet-50-2 has 2048-1024-2048. 198 | Args: 199 | pretrained (bool): If True, returns a model pre-trained on ImageNet 200 | progress (bool): If True, displays a progress bar of the download to stderr 201 | """ 202 | kwargs['width_per_group'] = 64 * 2 203 | return _custom_resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) 204 | 205 | 206 | def custom_wide_resnet101_2(pretrained=False, progress=True, **kwargs): 207 | r"""Wide ResNet-101-2 model from 208 | `"Wide Residual Networks" `_ 209 | The model is the same as ResNet except for the bottleneck number of channels 210 | which is twice larger in every block. The number of channels in outer 1x1 211 | convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 212 | channels, and in Wide ResNet-50-2 has 2048-1024-2048. 213 | Args: 214 | pretrained (bool): If True, returns a model pre-trained on ImageNet 215 | progress (bool): If True, displays a progress bar of the download to stderr 216 | """ 217 | kwargs['width_per_group'] = 64 * 2 218 | return _custom_resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) 219 | --------------------------------------------------------------------------------