Official implementation of paper "Avatar Knowledge Distillation: Self-ensemble Teacher Paradigm with Uncertainty", ACM MM 2023.
# dataset settings
dataset_type = 'CIFAR10'
img_norm_cfg = dict(
    mean=[125.307, 122.961, 113.8575],
    std=[51.5865, 50.847, 51.255],
    to_rgb=False)
train_pipeline = [
    dict(type='RandomCrop', size=32, padding=4),
    dict(type='RandomFlip', flip_prob=0.5, direction='horizontal'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Cutout', shape=16, pad_val=0, prob=1.0),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='ToTensor', keys=['gt_label']),
    dict(type='Collect', keys=['img', 'gt_label'])
]
test_pipeline = [
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='Collect', keys=['img'])
]
data = dict(
    samples_per_gpu=96,
    workers_per_gpu=2,
    train=dict(
        type=dataset_type, data_prefix='data/cifar10',
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        data_prefix='data/cifar10',
        pipeline=test_pipeline,
        test_mode=True),
    test=dict(
        type=dataset_type,
        data_prefix='data/cifar10',
        pipeline=test_pipeline,
        test_mode=True))

configs/_base_/datasets/mmcls/cifar10_bs96_cutout.py
Gumpest/AvatarKD
# dataset settings
dataset_type = 'ImageNet'
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='RandomResizedCrop',
        size=224,
        scale=(0.25, 1.0),
        backend='pillow'),
    dict(type='ColorJitter', brightness=0.4, contrast=0.4, saturation=0.4),
    dict(type='RandomFlip', flip_prob=0.5, direction='horizontal'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='ToTensor', keys=['gt_label']),
    dict(type='Collect', keys=['img', 'gt_label'])
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', size=(256, -1), backend='pillow'),
    dict(type='CenterCrop', crop_size=224),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='Collect', keys=['img'])
]
data = dict(
    samples_per_gpu=256,
    workers_per_gpu=4,
    train=dict(
        type=dataset_type,
        data_prefix='data/imagenet/train',
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        data_prefix='data/imagenet/val',
        ann_file='data/imagenet/meta/val.txt',
        pipeline=test_pipeline),
    test=dict(
        # replace `data/val` with `data/test` for standard test
        type=dataset_type,
        data_prefix='data/imagenet/val',
        ann_file='data/imagenet/meta/val.txt',
        pipeline=test_pipeline))
evaluation = dict(interval=1, metric='accuracy')

configs/_base_/datasets/mmcls/imagenet_bs256_autoslim.py
Gumpest/AvatarKD
# optimizer
paramwise_cfg = dict(bias_decay_mult=0.0, norm_decay_mult=0.0)
optimizer = dict(
    type='SGD',
    lr=0.5,
    momentum=0.9,
    weight_decay=4e-5,
    paramwise_cfg=paramwise_cfg)
optimizer_config = dict(grad_clip=None)
# learning policy
lr_config = dict(policy='poly', power=1.0, min_lr=0.0, by_epoch=False)
runner = dict(type='IterBasedRunner', max_iters=300000)

configs/_base_/schedules/mmcls/imagenet_bs1024_spos.py
Gumpest/AvatarKD
# optimizer
paramwise_cfg = dict(
    bias_decay_mult=0.0, norm_decay_mult=0.0, dwconv_decay_mult=0.0)
optimizer = dict(
    type='SGD',
    lr=0.5,
    momentum=0.9,
    nesterov=True,
    weight_decay=0.0001,
    paramwise_cfg=paramwise_cfg)

optimizer_config = None

# learning policy
lr_config = dict(policy='poly', power=1.0, min_lr=0.0, by_epoch=False)
runner = dict(type='EpochBasedRunner', max_epochs=300)

configs/_base_/schedules/mmcls/imagenet_bs2048_autoslim.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmdet/coco_detection.py',
    '../../_base_/schedules/mmdet/schedule_2x.py',
    '../../_base_/mmdet_runtime.py'
]

# model settings
student = dict(
    type='mmdet.RepPointsDetector',
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch',
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        start_level=1,
        add_extra_convs='on_input',
        num_outs=5,
        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True),
        init_cfg=dict(type='Pretrained', prefix='neck', checkpoint='token_ckpt/reppoints_moment_x101_fpn_dconv_c3-c5_gn-neck+head_2x_coco_20200329-f87da1ea.pth')),
    bbox_head=dict(
        type='RepPointsHead',
        num_classes=80,
        in_channels=256,
        feat_channels=256,
        point_feat_channels=256,
        stacked_convs=3,
        num_points=9,
        gradient_mul=0.1,
        point_strides=[8, 16, 32, 64, 128],
        point_base_scale=4,
        loss_cls=dict(
            type='FocalLoss',
            use_sigmoid=True,
            gamma=2.0,
            alpha=0.25,
            loss_weight=1.0),
        loss_bbox_init=dict(type='SmoothL1Loss', beta=0.11, loss_weight=0.5),
        loss_bbox_refine=dict(type='SmoothL1Loss', beta=0.11, loss_weight=1.0),
        transform_method='moment',
        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True),
        init_cfg=dict(type='Pretrained', prefix='bbox_head', checkpoint='token_ckpt/reppoints_moment_x101_fpn_dconv_c3-c5_gn-neck+head_2x_coco_20200329-f87da1ea.pth')),
    # training and testing settings
    train_cfg=dict(
        init=dict(
            assigner=dict(type='PointAssigner', scale=4, pos_num=1),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        refine=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.5,
                neg_iou_thr=0.4,
                min_pos_iou=0,
                ignore_iof_thr=-1),
            allowed_border=-1,
            pos_weight=-1,
            debug=False)),
    test_cfg=dict(
        nms_pre=1000,
        min_bbox_size=0,
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.5),
        max_per_img=100))

teacher = dict(
    type='mmdet.RepPointsDetector',
    init_cfg=dict(type='Pretrained', checkpoint='https://download.openmmlab.com/mmdetection/v2.0/reppoints/reppoints_moment_x101_fpn_dconv_c3-c5_gn-neck%2Bhead_2x_coco/reppoints_moment_x101_fpn_dconv_c3-c5_gn-neck%2Bhead_2x_coco_20200329-f87da1ea.pth'),
    backbone=dict(
        type='ResNeXt',
        depth=101,
        groups=32,
        base_width=4,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        style='pytorch',
        dcn=dict(type='DCN', deform_groups=1, fallback_on_stride=False),
        stage_with_dcn=(False, True, True, True)),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        start_level=1,
        add_extra_convs='on_input',
        num_outs=5,
        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True)),
    bbox_head=dict(
        type='RepPointsHead',
        num_classes=80,
        in_channels=256,
        feat_channels=256,
        point_feat_channels=256,
        stacked_convs=3,
        num_points=9,
        gradient_mul=0.1,
        point_strides=[8, 16, 32, 64, 128],
        point_base_scale=4,
        loss_cls=dict(
            type='FocalLoss',
            use_sigmoid=True,
            gamma=2.0,
            alpha=0.25,
            loss_weight=1.0),
        loss_bbox_init=dict(type='SmoothL1Loss', beta=0.11, loss_weight=0.5),
        loss_bbox_refine=dict(type='SmoothL1Loss', beta=0.11, loss_weight=1.0),
        transform_method='moment',
        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True)),
    # training and testing settings
    train_cfg=dict(
        init=dict(
            assigner=dict(type='PointAssigner', scale=4, pos_num=1),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        refine=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.5,
                neg_iou_thr=0.4,
                min_pos_iou=0,
                ignore_iof_thr=-1),
            allowed_border=-1,
            pos_weight=-1,
            debug=False)),
    test_cfg=dict(
        nms_pre=1000,
        min_bbox_size=0,
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.5),
        max_per_img=100))

# algorithm setting
algorithm = dict(
    type='GeneralDistill',
    architecture=dict(
        type='MMDetArchitecture',
        model=student,
    ),
    distiller=dict(
        type='SingleTeacherDistiller',
        teacher=teacher,
        teacher_trainable=False,
        components=[
            dict(
                student_module='neck',
                teacher_module='neck',
                losses=[
                    dict(
                        type='AKD_ChannelWiseDivergence',
                        name='loss_cwd_fpn',
                        loss_weight=7.5,
                    ),
                ],
                align_module=dict(
                    type='conv2d',
                    num_modules=5,
                    student_channels=256,
                    teacher_channels=256,
                )
            ),
            dict(
                student_module='bbox_head.reppoints_pts_refine_out',
                teacher_module='bbox_head.reppoints_pts_refine_out',
                losses=[
                    dict(
                        type='mmdet.SmoothL1Loss',
                        beta=0.11,
                        name='loss_kd_reg_refine_head',
                        loss_weight=0.2,
                    )
                ]
            ),
            dict(
                student_module='bbox_head.reppoints_pts_init_out',
                teacher_module='bbox_head.reppoints_pts_init_out',
                losses=[
                    dict(
                        type='mmdet.SmoothL1Loss',
                        beta=0.11,
                        name='loss_kd_reg_init_head',
                        loss_weight=0.1,
                    )
                ]
            ),
        ]),
)

find_unused_parameters = True

opt_paramwise_cfg = dict(
    custom_keys={
        'distiller.teacher': dict(lr_mult=0, decay_mult=0),
        'distiller.losses': dict(lr_mult=0, decay_mult=0),
    }
)
# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001, paramwise_cfg=opt_paramwise_cfg)


configs/distill/akd/akd_cwd_reppoints_x101-reppoints-r50_coco.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmdet/coco_detection.py',
    '../../_base_/schedules/mmdet/schedule_2x.py',
    '../../_base_/mmdet_runtime.py'
]

# model settings
student = dict(
    type='mmdet.RetinaNet',
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch',
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        start_level=1,
        add_extra_convs='on_input',
        num_outs=5),
    bbox_head=dict(
        type='RetinaHead',
        num_classes=80,
        in_channels=256,
        stacked_convs=4,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            octave_base_scale=4,
            scales_per_octave=3,
            ratios=[0.5, 1.0, 2.0],
            strides=[8, 16, 32, 64, 128]),
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[.0, .0, .0, .0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        loss_cls=dict(
            type='FocalLoss',
            use_sigmoid=True,
            gamma=2.0,
            alpha=0.25,
            loss_weight=1.0),
        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
    # model training and testing settings
    train_cfg=dict(
        assigner=dict(
            type='MaxIoUAssigner',
            pos_iou_thr=0.5,
            neg_iou_thr=0.4,
            min_pos_iou=0,
            ignore_iof_thr=-1),
        allowed_border=-1,
        pos_weight=-1,
        debug=False),
    test_cfg=dict(
        nms_pre=1000,
        min_bbox_size=0,
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.5),
        max_per_img=100))

teacher = dict(
    type='mmdet.RetinaNet',
    init_cfg=dict(type='Pretrained', checkpoint='https://download.openmmlab.com/mmdetection/v2.0/retinanet/retinanet_r101_fpn_2x_coco/retinanet_r101_fpn_2x_coco_20200131-5560aee8.pth'),
    backbone=dict(
        type='ResNet',
        depth=101,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch'),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        start_level=1,
        add_extra_convs='on_input',
        num_outs=5),
    bbox_head=dict(
        type='RetinaHead',
        num_classes=80,
        in_channels=256,
        stacked_convs=4,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            octave_base_scale=4,
            scales_per_octave=3,
            ratios=[0.5, 1.0, 2.0],
            strides=[8, 16, 32, 64, 128]),
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[.0, .0, .0, .0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        loss_cls=dict(
            type='FocalLoss',
            use_sigmoid=True,
            gamma=2.0,
            alpha=0.25,
            loss_weight=1.0),
        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
    # model training and testing settings
    train_cfg=dict(
        assigner=dict(
            type='MaxIoUAssigner',
            pos_iou_thr=0.5,
            neg_iou_thr=0.4,
            min_pos_iou=0,
            ignore_iof_thr=-1),
        allowed_border=-1,
        pos_weight=-1,
        debug=False),
    test_cfg=dict(
        nms_pre=1000,
        min_bbox_size=0,
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.5),
        max_per_img=100))

# algorithm setting
algorithm = dict(
    type='GeneralDistill',
    architecture=dict(
        type='MMDetArchitecture',
        model=student,
    ),
    distiller=dict(
        type='SingleTeacherDistiller',
        teacher=teacher,
        teacher_trainable=False,
        components=[
            dict(
                student_module='neck',
                teacher_module='neck',
                losses=[
                    dict(
                        type='AKD_MasKDLoss',
                        name='loss_maskd_fpn',
                        channels=[256, 256, 256, 256, 256],
                        num_tokens=6,
                        weight_mask=True,
                        custom_mask=True,
                        custom_mask_warmup=7000,
                        pretrained='https://github.com/Gumpest/MasKD/releases/download/v0.0.3/maskd_retinanet_r101_fpn_2x_coco_20200131-5560aee8.ckpt',
                        loss_weight=7,
                    ),
                ],
                align_module=dict(
                    type='conv2d',
                    num_modules=5,
                    student_channels=256,
                    teacher_channels=256,
                )
            ),
            dict(
                student_module='bbox_head.retina_reg',
                teacher_module='bbox_head.retina_reg',
                losses=[
                    dict(
                        type='mmdet.L1Loss',
                        name='loss_kd_reg_head',
                        loss_weight=1, 
                    )
                ]
            ),
        ]),
)

find_unused_parameters = True

opt_paramwise_cfg = dict(
    custom_keys={
        'distiller.teacher': dict(lr_mult=0, decay_mult=0),
        'distiller.losses': dict(lr_mult=0, decay_mult=0),
    }
)

# optimizer
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001, paramwise_cfg=opt_paramwise_cfg)
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=35, norm_type=2))
lr_config = dict(warmup_iters=1000)

configs/distill/akd/akd_maskd_retinanet_x101-retinanet_r50_coco.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmdet/coco_detection.py',
    '../../_base_/schedules/mmdet/schedule_2x.py',
    '../../_base_/mmdet_runtime.py'
]

# model settings
student = dict(
    type='mmdet.FasterRCNN',
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch',
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        num_outs=5,
        init_cfg=dict(type='Pretrained', prefix='neck', checkpoint='https://download.openmmlab.com/mmdetection/v2.0/cascade_rcnn/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco_20200512_161033-bdb5126a.pth')),
    rpn_head=dict(
        type='RPNHead',
        in_channels=256,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            scales=[8],
            ratios=[0.5, 1.0, 2.0],
            strides=[4, 8, 16, 32, 64]),
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[.0, .0, .0, .0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        loss_cls=dict(
            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
        loss_bbox=dict(type='L1Loss', loss_weight=1.0),
        init_cfg=dict(type='Pretrained', prefix='rpn_head', checkpoint='https://download.openmmlab.com/mmdetection/v2.0/cascade_rcnn/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco_20200512_161033-bdb5126a.pth')),
    roi_head=dict(
        type='StandardRoIHead',
        bbox_roi_extractor=dict(
            type='SingleRoIExtractor',
            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
            out_channels=256,
            featmap_strides=[4, 8, 16, 32]),
        bbox_head=dict(
            type='Shared2FCBBoxHead',
            in_channels=256,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=80,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=True,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
            loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
        # init_cfg=dict(type='Pretrained', prefix='roi_head.bbox_head.0', checkpoint='https://download.openmmlab.com/mmdetection/v2.0/cascade_rcnn/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco_20200512_161033-bdb5126a.pth')),
    # model training and testing settings
    train_cfg=dict(
        rpn=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.7,
                neg_iou_thr=0.3,
                min_pos_iou=0.3,
                match_low_quality=True,
                ignore_iof_thr=-1),
            sampler=dict(
                type='RandomSampler',
                num=256,
                pos_fraction=0.5,
                neg_pos_ub=-1,
                add_gt_as_proposals=False),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        rpn_proposal=dict(
            nms_pre=2000,
            max_per_img=1000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.5,
                neg_iou_thr=0.5,
                min_pos_iou=0.5,
                match_low_quality=False,
                ignore_iof_thr=-1),
            sampler=dict(
                type='RandomSampler',
                num=512,
                pos_fraction=0.25,
                neg_pos_ub=-1,
                add_gt_as_proposals=True),
            pos_weight=-1,
            debug=False)),
    test_cfg=dict(
        rpn=dict(
            nms_pre=1000,
            max_per_img=1000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=dict(
            score_thr=0.05,
            nms=dict(type='nms', iou_threshold=0.5),
            max_per_img=100)
        # soft-nms is also supported for rcnn testing
        # e.g., nms=dict(type='soft_nms', iou_threshold=0.5, min_score=0.05)
    ))


teacher = dict(
    type='mmdet.CascadeRCNN',
    init_cfg=dict(type='Pretrained', checkpoint='https://download.openmmlab.com/mmdetection/v2.0/cascade_rcnn/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco/cascade_mask_rcnn_x101_64x4d_fpn_20e_coco_20200512_161033-bdb5126a.pth'),
    backbone=dict(
        type='ResNeXt',
        depth=101,
        groups=64,
        base_width=4,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        style='pytorch'),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        num_outs=5),
    rpn_head=dict(
        type='RPNHead',
        in_channels=256,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            scales=[8],
            ratios=[0.5, 1.0, 2.0],
            strides=[4, 8, 16, 32, 64]),
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[.0, .0, .0, .0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        loss_cls=dict(
            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
        loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)),
    roi_head=dict(
        type='CascadeRoIHead',
        num_stages=1,
        stage_loss_weights=[1, 0.5, 0.25],
        bbox_roi_extractor=dict(
            type='SingleRoIExtractor',
            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
            out_channels=256,
            featmap_strides=[4, 8, 16, 32]),
        bbox_head=[
            dict(
                # NOTE: use `KDShared2FCBBoxHead` to discard loss computation
                type='KDShared2FCBBoxHead',
                in_channels=256,
                fc_out_channels=1024,
                roi_feat_size=7,
                num_classes=80,
                bbox_coder=dict(
                    type='DeltaXYWHBBoxCoder',
                    target_means=[0., 0., 0., 0.],
                    target_stds=[0.1, 0.1, 0.2, 0.2]),
                reg_class_agnostic=True,
                loss_cls=dict(
                    type='CrossEntropyLoss',
                    use_sigmoid=False,
                    loss_weight=1.0),
                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
                               loss_weight=1.0)),
            # dict(
            #     type='Shared2FCBBoxHead',
            #     in_channels=256,
            #     fc_out_channels=1024,
            #     roi_feat_size=7,
            #     num_classes=80,
            #     bbox_coder=dict(
            #         type='DeltaXYWHBBoxCoder',
            #         target_means=[0., 0., 0., 0.],
            #         target_stds=[0.05, 0.05, 0.1, 0.1]),
            #     reg_class_agnostic=True,
            #     loss_cls=dict(
            #         type='CrossEntropyLoss',
            #         use_sigmoid=False,
            #         loss_weight=1.0),
            #     loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
            #                    loss_weight=1.0)),
            # dict(
            #     type='Shared2FCBBoxHead',
            #     in_channels=256,
            #     fc_out_channels=1024,
            #     roi_feat_size=7,
            #     num_classes=80,
            #     bbox_coder=dict(
            #         type='DeltaXYWHBBoxCoder',
            #         target_means=[0., 0., 0., 0.],
            #         target_stds=[0.033, 0.033, 0.067, 0.067]),
            #     reg_class_agnostic=True,
            #     loss_cls=dict(
            #         type='CrossEntropyLoss',
            #         use_sigmoid=False,
            #         loss_weight=1.0),
            #     loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=1.0))
        ],
        # mask_roi_extractor=dict(
        #     type='SingleRoIExtractor',
        #     roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
        #     out_channels=256,
        #     featmap_strides=[4, 8, 16, 32]),
        # mask_head=dict(
        #     type='FCNMaskHead',
        #     num_convs=4,
        #     in_channels=256,
        #     conv_out_channels=256,
        #     num_classes=80,
        #     loss_mask=dict(
        #         type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))),
        ),
    # model training and testing settings
    train_cfg=dict(
        rpn=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.7,
                neg_iou_thr=0.3,
                min_pos_iou=0.3,
                match_low_quality=True,
                ignore_iof_thr=-1),
            sampler=dict(
                type='RandomSampler',
                num=256,
                pos_fraction=0.5,
                neg_pos_ub=-1,
                add_gt_as_proposals=False),
            allowed_border=0,
            pos_weight=-1,
            debug=False),
        rpn_proposal=dict(
            nms_pre=2000,
            max_per_img=2000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=[
            dict(
                assigner=dict(
                    type='MaxIoUAssigner',
                    pos_iou_thr=0.5,
                    neg_iou_thr=0.5,
                    min_pos_iou=0.5,
                    match_low_quality=False,
                    ignore_iof_thr=-1),
                sampler=dict(
                    type='RandomSampler',
                    num=512,
                    pos_fraction=0.25,
                    neg_pos_ub=-1,
                    add_gt_as_proposals=True),
                pos_weight=-1,
                debug=False),
            # dict(
            #     assigner=dict(
            #         type='MaxIoUAssigner',
            #         pos_iou_thr=0.6,
            #         neg_iou_thr=0.6,
            #         min_pos_iou=0.6,
            #         match_low_quality=False,
            #         ignore_iof_thr=-1),
            #     sampler=dict(
            #         type='RandomSampler',
            #         num=512,
            #         pos_fraction=0.25,
            #         neg_pos_ub=-1,
            #         add_gt_as_proposals=True),
            #     pos_weight=-1,
            #     debug=False),
            # dict(
            #     assigner=dict(
            #         type='MaxIoUAssigner',
            #         pos_iou_thr=0.7,
            #         neg_iou_thr=0.7,
            #         min_pos_iou=0.7,
            #         match_low_quality=False,
            #         ignore_iof_thr=-1),
            #     sampler=dict(
            #         type='RandomSampler',
            #         num=512,
            #         pos_fraction=0.25,
            #         neg_pos_ub=-1,
            #         add_gt_as_proposals=True),
            #     pos_weight=-1,
            #     debug=False)
        ]),
    test_cfg=dict(
        rpn=dict(
            nms_pre=1000,
            max_per_img=1000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=dict(
            score_thr=0.05,
            nms=dict(type='nms', iou_threshold=0.5),
            max_per_img=100,
            )))


# algorithm setting
algorithm = dict(
    type='GeneralDistill',
    architecture=dict(
        type='MMDetArchitecture',
        model=student,
    ),
    distiller=dict(
        type='SingleTeacherDistiller',
        teacher=teacher,
        teacher_trainable=False,
        components=[
            dict(
                student_module='neck',
                teacher_module='neck',
                losses=[
                    dict(
                        type='AKD_MGDLoss',
                        name='loss_mgd_fpn',
                        loss_weight=1,
                    ),
                ],
                align_module=dict(
                    type='conv2d',
                    num_modules=5,
                    student_channels=256,
                    teacher_channels=256,
                )
            ),
            dict(
                student_module='roi_head.bbox_head.fc_reg',
                teacher_module='roi_head.bbox_head.0.fc_reg',
                modules_with_student_inputs=[
                    dict(
                        student_module='roi_head.bbox_roi_extractor',
                        teacher_module='roi_head.bbox_roi_extractor.0',
                        same_indices=[1],)
                ],
                losses=[
                    dict(
                        type='mmdet.SmoothL1Loss',
                        name='loss_kd_roi_reg_head',
                        loss_weight=1,
                    )
                ]
            ),
            dict(
                student_module='roi_head.bbox_head.fc_cls',
                teacher_module='roi_head.bbox_head.0.fc_cls',
                modules_with_student_inputs=[
                    dict(
                        student_module='roi_head.bbox_roi_extractor',
                        teacher_module='roi_head.bbox_roi_extractor.0',
                        same_indices=[1],)
                ],
                losses=[
                    dict(
                        type='DIST',
                        name='loss_dist_roi_cls_head',
                        beta=1,
                        gamma=1,
                        loss_weight=1,
                    )
                ]),
        ]),
)

find_unused_parameters = True

opt_paramwise_cfg = dict(
    custom_keys={
        'distiller.teacher': dict(lr_mult=0, decay_mult=0),
        'distiller.losses': dict(lr_mult=0, decay_mult=0),
    }
)
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=35, norm_type=2))
optimizer = dict(paramwise_cfg=opt_paramwise_cfg)
configs/distill/akd/akd_mgd_cascade_mask_rcnn_x101-fpn_x50_coco.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmdet/coco_detection.py',
    '../../_base_/schedules/mmdet/schedule_1x.py',
    '../../_base_/mmdet_runtime.py'
]

# model settings
student = dict(
    type='mmdet.GFL',
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch',
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        start_level=1,
        add_extra_convs='on_output',
        num_outs=5),
    bbox_head=dict(
        type='GFLHead',
        num_classes=80,
        in_channels=256,
        stacked_convs=4,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            ratios=[1.0],
            octave_base_scale=8,
            scales_per_octave=1,
            strides=[8, 16, 32, 64, 128]),
        loss_cls=dict(
            type='QualityFocalLoss',
            use_sigmoid=True,
            beta=2.0,
            loss_weight=1.0),
        loss_dfl=dict(type='DistributionFocalLoss', loss_weight=0.25),
        reg_max=16,
        loss_bbox=dict(type='GIoULoss', loss_weight=2.0)),
    # training and testing settings
    train_cfg=dict(
        assigner=dict(type='ATSSAssigner', topk=9),
        allowed_border=-1,
        pos_weight=-1,
        debug=False),
    test_cfg=dict(
        nms_pre=1000,
        min_bbox_size=0,
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.6),
        max_per_img=100))

checkpoint = 'https://download.openmmlab.com/mmdetection/v2.0/gfl/gfl_r101_fpn_mstrain_2x_coco/gfl_r101_fpn_mstrain_2x_coco_20200629_200126-dd12f847.pth'  # noqa: E501

teacher = dict(
    type='mmdet.GFL',
    init_cfg=dict(type='Pretrained', checkpoint=checkpoint),
    backbone=dict(
        type='ResNet',
        depth=101,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch',
        init_cfg=None),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        start_level=1,
        add_extra_convs='on_output',
        num_outs=5),
    bbox_head=dict(
        type='GFLHead',
        num_classes=80,
        in_channels=256,
        stacked_convs=4,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            ratios=[1.0],
            octave_base_scale=8,
            scales_per_octave=1,
            strides=[8, 16, 32, 64, 128]),
        loss_cls=dict(
            type='QualityFocalLoss',
            use_sigmoid=True,
            beta=2.0,
            loss_weight=1.0),
        loss_dfl=dict(type='DistributionFocalLoss', loss_weight=0.25),
        reg_max=16,
        loss_bbox=dict(type='GIoULoss', loss_weight=2.0)),
    # training and testing settings
    train_cfg=dict(
        assigner=dict(type='ATSSAssigner', topk=9),
        allowed_border=-1,
        pos_weight=-1,
        debug=False),
    test_cfg=dict(
        nms_pre=1000,
        min_bbox_size=0,
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.6),
        max_per_img=100))

# algorithm setting
algorithm = dict(
    type='GeneralDistill',
    architecture=dict(
        type='MMDetArchitecture',
        model=student,
    ),
    distiller=dict(
        type='SingleTeacherDistiller',
        teacher=teacher,
        teacher_trainable=False,
        components=[
            dict(
                student_module='bbox_head.gfl_cls',
                teacher_module='bbox_head.gfl_cls',
                losses=[
                    dict(
                        type='ChannelWiseDivergence',
                        name='loss_cwd_cls_head',
                        tau=1,
                        loss_weight=5,
                    )
                ])
        ]),
)

find_unused_parameters = True

# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)

configs/distill/cwd/cwd_cls_head_gfl_r101_fpn_gfl_r50_fpn_1x_coco.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmseg/cityscapes.py',
    '../../_base_/mmseg_runtime.py',
    '../../_base_/schedules/mmseg/schedule_80k.py'
]

norm_cfg = dict(type='SyncBN', requires_grad=True)

# pspnet r18
student = dict(
    type='mmseg.EncoderDecoder',
    backbone=dict(
        type='ResNetV1c',
        init_cfg=dict(
            type='Pretrained', checkpoint='open-mmlab://resnet18_v1c'),
        depth=18,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        dilations=(1, 1, 2, 4),
        strides=(1, 2, 1, 1),
        norm_cfg=norm_cfg,
        norm_eval=False,
        style='pytorch',
        contract_dilation=True),
    decode_head=dict(
        type='PSPHead',
        in_channels=512,
        in_index=3,
        channels=128,
        pool_scales=(1, 2, 3, 6),
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=256,
        in_index=2,
        channels=64,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pspnet/pspnet_r101-d8_512x1024_80k_cityscapes/pspnet_r101-d8_512x1024_80k_cityscapes_20200606_112211-e1e1100f.pth'  # noqa: E501

# pspnet r101
teacher = dict(
    type='mmseg.EncoderDecoder',
    init_cfg=dict(type='Pretrained', checkpoint=checkpoint),
    backbone=dict(
        type='ResNetV1c',
        depth=101,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        dilations=(1, 1, 2, 4),
        strides=(1, 2, 1, 1),
        norm_cfg=norm_cfg,
        norm_eval=False,
        style='pytorch',
        contract_dilation=True),
    decode_head=dict(
        type='PSPHead',
        in_channels=2048,
        in_index=3,
        channels=512,
        pool_scales=(1, 2, 3, 6),
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
)

# algorithm setting
algorithm = dict(
    type='GeneralDistill',
    architecture=dict(
        type='MMSegArchitecture',
        model=student,
    ),
    distiller=dict(
        type='SingleTeacherDistiller',
        teacher=teacher,
        teacher_trainable=False,
        components=[
            dict(
                student_module='decode_head.conv_seg',
                teacher_module='decode_head.conv_seg',
                losses=[
                    dict(
                        type='ChannelWiseDivergence',
                        name='loss_cwd_logits',
                        tau=1,
                        loss_weight=5,
                    )
                ])
        ]),
)

find_unused_parameters = True

configs/distill/cwd/cwd_cls_head_pspnet_r101_d8_pspnet_r18_d8_512x1024_cityscapes_80k.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmcls/imagenet_bs32.py',
    '../../_base_/schedules/mmcls/imagenet_bs256.py',
    '../../_base_/mmcls_runtime.py'
]

# model settings
student = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(
        type='ResNet',
        depth=18,
        num_stages=4,
        out_indices=(3, ),
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=512,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5),
    ))

# teacher settings
teacher_ckpt = 'https://download.openmmlab.com/mmclassification/v0/resnet/resnet34_8xb32_in1k_20210831-f257d4e6.pth'  # noqa: E501

teacher = dict(
    type='mmcls.ImageClassifier',
    init_cfg=dict(type='Pretrained', checkpoint=teacher_ckpt),
    backbone=dict(
        type='ResNet',
        depth=34,
        num_stages=4,
        out_indices=(3, ),
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=512,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5),
    ))

# algorithm setting
algorithm = dict(
    type='GeneralDistill',
    architecture=dict(
        type='MMClsArchitecture',
        model=student,
    ),
    with_student_loss=True,
    with_teacher_loss=False,
    distiller=dict(
        type='SingleTeacherDistiller',
        teacher=teacher,
        teacher_trainable=False,
        teacher_norm_eval=True,
        components=[
            dict(
                student_module='neck.gap',
                teacher_module='neck.gap',
                losses=[
                    dict(
                        type='DistanceWiseRKD',
                        name='distance_wise_loss',
                        loss_weight=25.0,
                        with_l2_norm=True),
                    dict(
                        type='AngleWiseRKD',
                        name='angle_wise_loss',
                        loss_weight=50.0,
                        with_l2_norm=True),
                ])
        ]),
)

find_unused_parameters = True

configs/distill/rkd/rkd_neck_resnet34_resnet18_8xb32_in1k.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmcls/imagenet_bs32.py',
    '../../_base_/schedules/mmcls/imagenet_bs256.py',
    '../../_base_/mmcls_runtime.py'
]

# model settings
student = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(
        type='ResNet',
        depth=18,
        num_stages=4,
        out_indices=(3, ),
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=512,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5),
    ))

checkpoint = 'https://download.openmmlab.com/mmclassification/v0/resnet/resnet50_8xb32_in1k_20210831-ea4938fc.pth'  # noqa: E501

# teacher settings
teacher = dict(
    type='mmcls.ImageClassifier',
    init_cfg=dict(type='Pretrained', checkpoint=checkpoint),
    backbone=dict(
        type='ResNet',
        depth=34,
        num_stages=4,
        out_indices=(3, ),
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=512,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5),
    ))

# algorithm setting
algorithm = dict(
    type='GeneralDistill',
    architecture=dict(
        type='MMClsArchitecture',
        model=student,
    ),
    with_student_loss=True,
    # teacher_trainable and with_teacher_loss have a dependency
    # relationship, if teacher_trainable is false, then
    # with_teacher_loss must be false.
    with_teacher_loss=False,
    distiller=dict(
        type='SingleTeacherDistiller',
        teacher=teacher,
        teacher_trainable=False,
        teacher_norm_eval=True,
        components=[
            dict(
                student_module='head.fc',
                teacher_module='head.fc',
                losses=[
                    dict(
                        type='WSLD',
                        name='loss_wsld',
                        tau=2,
                        loss_weight=2.5,
                        num_classes=1000)
                ])
        ]),
)

find_unused_parameters = True

configs/distill/wsld/wsld_cls_head_resnet34_resnet18_8xb32_in1k.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmcls/cifar10_bs96_cutout.py',
    '../../_base_/mmcls_runtime.py'
]

model = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(
        type='DartsBackbone',
        in_channels=3,
        base_channels=36,
        num_layers=20,
        num_nodes=4,
        stem_multiplier=3,
        auxliary=True,
        aux_channels=128,
        aux_out_channels=768,
        out_indices=(19, )),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='DartsSubnetClsHead',
        num_classes=10,
        in_channels=576,
        aux_in_channels=768,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        aux_loss=dict(type='CrossEntropyLoss', loss_weight=0.4),
        topk=(1, 5),
        cal_acc=True),
)

# FIXME: you may replace this with the mutable_cfg searched by yourself
mutable_cfg = 'https://download.openmmlab.com/mmrazor/v0.1/nas/darts/darts_subnetnet_1xb96_cifar10/darts_subnetnet_1xb96_cifar10_acc-97.32_20211222-e5727921_mutable_cfg.yaml'  # noqa: E501

algorithm = dict(
    type='Darts',
    architecture=dict(type='MMClsArchitecture', model=model),
    mutator=dict(
        type='DartsMutator',
        placeholder_mapping=dict(
            node=dict(
                type='DifferentiableOP',
                with_arch_param=False,
                choices=dict(
                    zero=dict(type='DartsZero'),
                    skip_connect=dict(
                        type='DartsSkipConnect', use_drop_path=True),
                    max_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='max',
                        use_drop_path=True),
                    avg_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='avg',
                        use_drop_path=True),
                    sep_conv_3x3=dict(
                        type='DartsSepConv', kernel_size=3,
                        use_drop_path=True),
                    sep_conv_5x5=dict(
                        type='DartsSepConv', kernel_size=5,
                        use_drop_path=True),
                    dil_conv_3x3=dict(
                        type='DartsDilConv', kernel_size=3,
                        use_drop_path=True),
                    dil_conv_5x5=dict(
                        type='DartsDilConv', kernel_size=5,
                        use_drop_path=True),
                )),
            node_edge=dict(
                type='DifferentiableEdge',
                num_chosen=2,
                with_arch_param=False,
            )),
    ),
    retraining=True,
    unroll=False,
    mutable_cfg=mutable_cfg)

data = dict(workers_per_gpu=8)

# optimizer
optimizer = dict(type='SGD', lr=0.025, momentum=0.9, weight_decay=0.0003)
optimizer_config = dict(grad_clip=dict(max_norm=5, norm_type=2))
# learning policy
lr_config = dict(policy='CosineAnnealing', min_lr=0)
runner = dict(type='EpochBasedRunner', max_epochs=600)

custom_hooks = [
    dict(
        type='DropPathProbHook',
        max_prob=0.2,
        interval=1,
        priority=70,
    )
]

configs/nas/darts/darts_subnet_1xb96_cifar10.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmcls/cifar10_bs16.py',
    '../../_base_/mmcls_runtime.py'
]

model = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(
        type='DartsBackbone',
        in_channels=3,
        base_channels=16,
        num_layers=8,
        num_nodes=4,
        stem_multiplier=3,
        out_indices=(7, )),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=10,
        in_channels=256,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5),
        cal_acc=True),
)

algorithm = dict(
    type='Darts',
    architecture=dict(type='MMClsArchitecture', model=model),
    mutator=dict(
        type='DartsMutator',
        placeholder_mapping=dict(
            node=dict(
                type='DifferentiableOP',
                with_arch_param=True,
                choices=dict(
                    zero=dict(type='DartsZero'),
                    skip_connect=dict(
                        type='DartsSkipConnect',
                        norm_cfg=dict(type='BN', affine=False)),
                    max_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='max',
                        norm_cfg=dict(type='BN', affine=False)),
                    avg_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='avg',
                        norm_cfg=dict(type='BN', affine=False)),
                    sep_conv_3x3=dict(
                        type='DartsSepConv',
                        kernel_size=3,
                        norm_cfg=dict(type='BN', affine=False)),
                    sep_conv_5x5=dict(
                        type='DartsSepConv',
                        kernel_size=5,
                        norm_cfg=dict(type='BN', affine=False)),
                    dil_conv_3x3=dict(
                        type='DartsDilConv',
                        kernel_size=3,
                        norm_cfg=dict(type='BN', affine=False)),
                    dil_conv_5x5=dict(
                        type='DartsDilConv',
                        kernel_size=5,
                        norm_cfg=dict(type='BN', affine=False)),
                )),
            node_edge=dict(
                type='DifferentiableEdge',
                num_chosen=2,
                with_arch_param=False,
            )),
    ),
    retraining=False,
    unroll=True)

data = dict(samples_per_gpu=64, workers_per_gpu=8, split=True)

# optimizer
optimizer = dict(
    architecture=dict(type='SGD', lr=0.025, momentum=0.9, weight_decay=3e-4),
    mutator=dict(type='Adam', lr=3e-4, weight_decay=1e-3))
optimizer_config = None

# learning policy
lr_config = dict(
    policy='CosineAnnealing', min_lr=1e-3, freeze_optimizers=['mutator'])
runner = dict(type='MultiLoaderEpochBasedRunner', max_epochs=50)

custom_hooks = [dict(
    type='SearchSubnetHook',
    interval=1,
    priority=70,
)]

find_unused_parameter = True
use_ddp_wrapper = True

configs/nas/darts/darts_supernet_unroll_1xb64_cifar10.py
Gumpest/AvatarKD
_base_ = ['./detnas_supernet_frcnn_shufflenetv2_fpn_1x_coco.py']

data = dict(
    samples_per_gpu=128,
    workers_per_gpu=8,
)

algorithm = dict(bn_training_mode=True)

searcher = dict(
    type='EvolutionSearcher',
    metrics='bbox',
    score_key='bbox_mAP',
    constraints=dict(flops=300 * 1e6),
    candidate_pool_size=50,
    candidate_top_k=10,
    max_epoch=20,
    num_mutation=20,
    num_crossover=20,
)

configs/nas/detnas/detnas_evolution_search_frcnn_shufflenetv2_fpn_coco.py
Gumpest/AvatarKD
_base_ = ['./detnas_supernet_frcnn_shufflenetv2_fpn_1x_coco.py']

# FIXME: you may replace this with the mutable_cfg searched by yourself
mutable_cfg = 'https://download.openmmlab.com/mmrazor/v0.1/nas/detnas/detnas_subnet_frcnn_shufflenetv2_fpn_1x_coco/detnas_subnet_frcnn_shufflenetv2_fpn_1x_coco_bbox_backbone_flops-0.34M_mAP-37.5_20211222-67fea61f_mutable_cfg.yaml'  # noqa: E501

algorithm = dict(retraining=True, mutable_cfg=mutable_cfg)

configs/nas/detnas/detnas_subnet_frcnn_shufflenetv2_fpn_1x_coco.py
Gumpest/AvatarKD
_base_ = [
    '../spos/spos_subnet_shufflenetv2_8xb128_in1k.py',
]

# FIXME: you may replace this with the mutable_cfg searched by yourself
mutable_cfg = 'https://download.openmmlab.com/mmrazor/v0.1/nas/detnas/detnas_subnet_frcnn_shufflenetv2_fpn_1x_coco/detnas_subnet_frcnn_shufflenetv2_fpn_1x_coco_bbox_backbone_flops-0.34M_mAP-37.5_20211222-67fea61f_mutable_cfg.yaml'  # noqa: E501

algorithm = dict(mutable_cfg=mutable_cfg)

configs/nas/detnas/detnas_subnet_shufflenetv2_8xb128_in1k.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmdet/coco_detection.py',
    '../../_base_/schedules/mmdet/schedule_1x.py',
    '../../_base_/mmdet_runtime.py'
]

norm_cfg = dict(type='SyncBN', requires_grad=True)
model = dict(
    type='mmdet.FasterRCNN',
    backbone=dict(
        type='mmcls.SearchableShuffleNetV2',
        norm_cfg=norm_cfg,
        out_indices=(0, 1, 2, 3),
        widen_factor=1.0,
        with_last_layer=False),
    neck=dict(
        type='FPN',
        norm_cfg=norm_cfg,
        in_channels=[64, 160, 320, 640],
        out_channels=256,
        num_outs=5),
    rpn_head=dict(
        type='RPNHead',
        in_channels=256,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            scales=[8],
            ratios=[0.5, 1.0, 2.0],
            strides=[4, 8, 16, 32, 64]),
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[.0, .0, .0, .0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        loss_cls=dict(
            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
    roi_head=dict(
        type='StandardRoIHead',
        bbox_roi_extractor=dict(
            type='SingleRoIExtractor',
            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
            out_channels=256,
            featmap_strides=[4, 8, 16, 32]),
        bbox_head=dict(
            type='Shared4Conv1FCBBoxHead',
            norm_cfg=norm_cfg,
            in_channels=256,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=80,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=False,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
            loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
    train_cfg=dict(
        rpn=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.7,
                neg_iou_thr=0.3,
                min_pos_iou=0.3,
                match_low_quality=True,
                ignore_iof_thr=-1),
            sampler=dict(
                type='RandomSampler',
                num=256,
                pos_fraction=0.5,
                neg_pos_ub=-1,
                add_gt_as_proposals=False),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        rpn_proposal=dict(
            nms_pre=2000,
            max_per_img=1000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.5,
                neg_iou_thr=0.5,
                min_pos_iou=0.5,
                match_low_quality=False,
                ignore_iof_thr=-1),
            sampler=dict(
                type='RandomSampler',
                num=512,
                pos_fraction=0.25,
                neg_pos_ub=-1,
                add_gt_as_proposals=True),
            pos_weight=-1,
            debug=False)),
    test_cfg=dict(
        rpn=dict(
            nms_pre=1000,
            max_per_img=1000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=dict(
            score_thr=0.05,
            nms=dict(type='nms', iou_threshold=0.5),
            max_per_img=100)
        # soft-nms is also supported for rcnn testing
        # e.g., nms=dict(type='soft_nms', iou_threshold=0.5, min_score=0.05)
    ),
)

mutator = dict(
    type='OneShotMutator',
    placeholder_mapping=dict(
        all_blocks=dict(
            type='OneShotOP',
            choices=dict(
                shuffle_3x3=dict(
                    type='ShuffleBlock', norm_cfg=norm_cfg, kernel_size=3),
                shuffle_5x5=dict(
                    type='ShuffleBlock', norm_cfg=norm_cfg, kernel_size=5),
                shuffle_7x7=dict(
                    type='ShuffleBlock', norm_cfg=norm_cfg, kernel_size=7),
                shuffle_xception=dict(
                    type='ShuffleXception',
                    norm_cfg=norm_cfg,
                ),
            ))))

algorithm = dict(
    type='DetNAS',
    architecture=dict(
        type='MMDetArchitecture',
        model=model,
    ),
    mutator=mutator,
    pruner=None,
    distiller=None,
    retraining=False,
)

find_unused_parameters = True

configs/nas/detnas/detnas_supernet_frcnn_shufflenetv2_fpn_1x_coco.py
Gumpest/AvatarKD
_base_ = [
    '../spos/spos_supernet_shufflenetv2_8xb128_in1k.py',
]

runner = dict(max_iters=300000)

configs/nas/detnas/detnas_supernet_shufflenetv2_8xb128_in1k.py
Gumpest/AvatarKD
_base_ = ['./spos_supernet_mobilenet_proxyless_gpu_8xb128_in1k.py']

data = dict(
    samples_per_gpu=512,
    workers_per_gpu=16,
)

algorithm = dict(bn_training_mode=True)

searcher = dict(
    type='EvolutionSearcher',
    candidate_pool_size=50,
    candidate_top_k=10,
    constraints=dict(flops=465 * 1e6),
    metrics='accuracy',
    score_key='accuracy_top-1',
    max_epoch=20,
    num_mutation=25,
    num_crossover=25,
    mutate_prob=0.1)

configs/nas/spos/spos_evolution_search_mobilenet_proxyless_gpu_flops465_8xb512_in1k.py
Gumpest/AvatarKD
_base_ = ['./spos_supernet_shufflenetv2_8xb128_in1k.py']

data = dict(
    samples_per_gpu=2048,
    workers_per_gpu=16,
)

algorithm = dict(bn_training_mode=True)

searcher = dict(
    type='EvolutionSearcher',
    candidate_pool_size=50,
    candidate_top_k=10,
    constraints=dict(flops=330 * 1e6),
    metrics='accuracy',
    score_key='accuracy_top-1',
    max_epoch=20,
    num_mutation=25,
    num_crossover=25,
    mutate_prob=0.1)

configs/nas/spos/spos_evolution_search_shufflenetv2_8xb2048_in1k.py
Gumpest/AvatarKD
_base_ = [
    './spos_subnet_mobilenet_proxyless_gpu_8xb128_in1k.py',
]

img_norm_cfg = dict(mean=[0., 0., 0.], std=[1., 1., 1.], to_rgb=False)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='RandomResizedCrop', size=224),
    dict(type='ColorJitter', brightness=0.4, contrast=0.4, saturation=0.4),
    dict(type='RandomFlip', flip_prob=0.5, direction='horizontal'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='ToTensor', keys=['gt_label']),
    dict(type='Collect', keys=['img', 'gt_label'])
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', size=(256, -1)),
    dict(type='CenterCrop', crop_size=224),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='Collect', keys=['img'])
]
data = dict(
    train=dict(pipeline=train_pipeline),
    val=dict(pipeline=test_pipeline),
    test=dict(pipeline=test_pipeline))

configs/nas/spos/spos_mobilenet_for_check_ckpt_from_anglenas.py
Gumpest/AvatarKD
_base_ = [
    './spos_supernet_mobilenet_proxyless_gpu_8xb128_in1k.py',
]

# FIXME: you may replace this with the mutable_cfg searched by yourself
mutable_cfg = 'https://download.openmmlab.com/mmrazor/v0.1/nas/spos/spos_mobilenet_subnet/spos_angelnas_flops_0.49G_acc_75.98_20220307-54f4698f_mutable_cfg.yaml'  # noqa: E501

algorithm = dict(retraining=True, mutable_cfg=mutable_cfg)
evaluation = dict(interval=10000, metric='accuracy')
checkpoint_config = dict(interval=30000)

runner = dict(max_iters=300000)
find_unused_parameters = False

configs/nas/spos/spos_subnet_mobilenet_proxyless_gpu_8xb128_in1k.py
Gumpest/AvatarKD
_base_ = [
    './spos_supernet_shufflenetv2_8xb128_in1k.py',
]

# FIXME: you may replace this with the mutable_cfg searched by yourself
mutable_cfg = 'https://download.openmmlab.com/mmrazor/v0.1/nas/spos/spos_shufflenetv2_subnet_8xb128_in1k/spos_shufflenetv2_subnet_8xb128_in1k_flops_0.33M_acc_73.87_20211222-454627be_mutable_cfg.yaml'  # noqa: E501

algorithm = dict(retraining=True, mutable_cfg=mutable_cfg)

runner = dict(max_iters=300000)
find_unused_parameters = False

configs/nas/spos/spos_subnet_shufflenetv2_8xb128_in1k.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmcls/imagenet_bs128_colorjittor.py',
    '../../_base_/schedules/mmcls/imagenet_bs1024_spos.py',
    '../../_base_/mmcls_runtime.py'
]
norm_cfg = dict(type='BN')
model = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(
        type='SearchableMobileNet',
        first_channels=40,
        last_channels=1728,
        widen_factor=1.0,
        norm_cfg=norm_cfg,
        arch_setting_type='proxyless_gpu'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=1728,
        loss=dict(
            type='LabelSmoothLoss',
            num_classes=1000,
            label_smooth_val=0.1,
            mode='original',
            loss_weight=1.0),
        topk=(1, 5),
    ),
)

mutator = dict(
    type='OneShotMutator',
    placeholder_mapping=dict(
        searchable_blocks=dict(
            type='OneShotOP',
            choices=dict(
                mb_k3e3=dict(
                    type='MBBlock',
                    kernel_size=3,
                    expand_ratio=3,
                    norm_cfg=norm_cfg,
                    act_cfg=dict(type='ReLU6')),
                mb_k5e3=dict(
                    type='MBBlock',
                    kernel_size=5,
                    expand_ratio=3,
                    norm_cfg=norm_cfg,
                    act_cfg=dict(type='ReLU6')),
                mb_k7e3=dict(
                    type='MBBlock',
                    kernel_size=7,
                    expand_ratio=3,
                    norm_cfg=norm_cfg,
                    act_cfg=dict(type='ReLU6')),
                mb_k3e6=dict(
                    type='MBBlock',
                    kernel_size=3,
                    expand_ratio=6,
                    norm_cfg=norm_cfg,
                    act_cfg=dict(type='ReLU6')),
                mb_k5e6=dict(
                    type='MBBlock',
                    kernel_size=5,
                    expand_ratio=6,
                    norm_cfg=norm_cfg,
                    act_cfg=dict(type='ReLU6')),
                mb_k7e6=dict(
                    type='MBBlock',
                    kernel_size=7,
                    expand_ratio=6,
                    norm_cfg=norm_cfg,
                    act_cfg=dict(type='ReLU6')),
                identity=dict(type='Identity'))),
        first_blocks=dict(
            type='OneShotOP',
            choices=dict(
                mb_k3e1=dict(
                    type='MBBlock',
                    kernel_size=3,
                    expand_ratio=1,
                    norm_cfg=norm_cfg,
                    act_cfg=dict(type='ReLU6')), ))))

algorithm = dict(
    type='SPOS',
    architecture=dict(
        type='MMClsArchitecture',
        model=model,
    ),
    mutator=mutator,
    distiller=None,
    retraining=False,
)

runner = dict(max_iters=150000)
evaluation = dict(interval=10000, metric='accuracy')

# checkpoint saving
checkpoint_config = dict(interval=30000)

find_unused_parameters = True

configs/nas/spos/spos_supernet_mobilenet_proxyless_gpu_8xb128_in1k.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmcls/imagenet_bs128_colorjittor.py',
    '../../_base_/schedules/mmcls/imagenet_bs1024_spos.py',
    '../../_base_/mmcls_runtime.py'
]
norm_cfg = dict(type='BN')
model = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(
        type='SearchableShuffleNetV2', widen_factor=1.0, norm_cfg=norm_cfg),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=1024,
        loss=dict(
            type='LabelSmoothLoss',
            num_classes=1000,
            label_smooth_val=0.1,
            mode='original',
            loss_weight=1.0),
        topk=(1, 5),
    ),
)

mutator = dict(
    type='OneShotMutator',
    placeholder_mapping=dict(
        all_blocks=dict(
            type='OneShotOP',
            choices=dict(
                shuffle_3x3=dict(
                    type='ShuffleBlock', kernel_size=3, norm_cfg=norm_cfg),
                shuffle_5x5=dict(
                    type='ShuffleBlock', kernel_size=5, norm_cfg=norm_cfg),
                shuffle_7x7=dict(
                    type='ShuffleBlock', kernel_size=7, norm_cfg=norm_cfg),
                shuffle_xception=dict(
                    type='ShuffleXception', norm_cfg=norm_cfg),
            ))))

algorithm = dict(
    type='SPOS',
    architecture=dict(
        type='MMClsArchitecture',
        model=model,
    ),
    mutator=mutator,
    distiller=None,
    retraining=False,
)

runner = dict(max_iters=150000)
evaluation = dict(interval=1000, metric='accuracy')

# checkpoint saving
checkpoint_config = dict(interval=1000)

find_unused_parameters = True

configs/nas/spos/spos_supernet_shufflenetv2_8xb128_in1k.py
Gumpest/AvatarKD
_base_ = [
    './autoslim_mbv2_supernet_8xb256_in1k.py',
]

algorithm = dict(distiller=None, input_shape=(3, 224, 224))

searcher = dict(
    type='GreedySearcher',
    target_flops=[500000000, 300000000, 200000000],
    max_channel_bins=12,
    metrics='accuracy')

data = dict(samples_per_gpu=1024, workers_per_gpu=4)

configs/pruning/autoslim/autoslim_mbv2_search_8xb1024_in1k.py
Gumpest/AvatarKD
_base_ = [
    './autoslim_mbv2_supernet_8xb256_in1k.py',
]

model = dict(
    head=dict(
        loss=dict(
            type='LabelSmoothLoss',
            mode='original',
            label_smooth_val=0.1,
            loss_weight=1.0)))

# FIXME: you may replace this with the channel_cfg searched by yourself
channel_cfg = [
    'https://download.openmmlab.com/mmrazor/v0.1/pruning/autoslim/autoslim_mbv2_subnet_8xb256_in1k/autoslim_mbv2_subnet_8xb256_in1k_flops-0.53M_acc-74.23_20211222-e5208bbd_channel_cfg.yaml',  # noqa: E501
    'https://download.openmmlab.com/mmrazor/v0.1/pruning/autoslim/autoslim_mbv2_subnet_8xb256_in1k/autoslim_mbv2_subnet_8xb256_in1k_flops-0.32M_acc-72.73_20211222-b5b0b33c_channel_cfg.yaml',  # noqa: E501
    'https://download.openmmlab.com/mmrazor/v0.1/pruning/autoslim/autoslim_mbv2_subnet_8xb256_in1k/autoslim_mbv2_subnet_8xb256_in1k_flops-0.22M_acc-71.39_20211222-43117c7b_channel_cfg.yaml'  # noqa: E501
]

algorithm = dict(
    architecture=dict(type='MMClsArchitecture', model=model),
    distiller=None,
    retraining=True,
    bn_training_mode=False,
    channel_cfg=channel_cfg)

runner = dict(type='EpochBasedRunner', max_epochs=300)

find_unused_parameters = True

configs/pruning/autoslim/autoslim_mbv2_subnet_8xb256_in1k.py
Gumpest/AvatarKD
_base_ = [
    '../../_base_/datasets/mmcls/imagenet_bs256_autoslim.py',
    '../../_base_/schedules/mmcls/imagenet_bs2048_autoslim.py',
    '../../_base_/mmcls_runtime.py'
]

model = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(type='MobileNetV2', widen_factor=1.5),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=1920,
        loss=dict(
            type='LabelSmoothLoss',
            mode='original',
            label_smooth_val=0.1,
            loss_weight=1.0),
        topk=(1, 5),
    ))

algorithm = dict(
    type='AutoSlim',
    architecture=dict(type='MMClsArchitecture', model=model),
    distiller=dict(
        type='SelfDistiller',
        components=[
            dict(
                student_module='head.fc',
                teacher_module='head.fc',
                losses=[
                    dict(
                        type='KLDivergence',
                        name='loss_kd',
                        tau=1,
                        loss_weight=1,
                    )
                ]),
        ]),
    pruner=dict(
        type='RatioPruner',
        ratios=(2 / 12, 3 / 12, 4 / 12, 5 / 12, 6 / 12, 7 / 12, 8 / 12, 9 / 12,
                10 / 12, 11 / 12, 1.0)),
    retraining=False,
    bn_training_mode=True,
    input_shape=None)

runner = dict(type='EpochBasedRunner', max_epochs=50)

use_ddp_wrapper = True

configs/pruning/autoslim/autoslim_mbv2_supernet_8xb256_in1k.py
Gumpest/AvatarKD
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import os
import subprocess  # noqa: F401
import sys

import pytorch_sphinx_theme

sys.path.insert(0, os.path.abspath('../..'))

# -- Project information -----------------------------------------------------

project = 'MMRazor'
copyright = '2021-2030, OpenMMLab'
author = 'MMRazor Authors'
version_file = '../../mmrazor/version.py'


def get_version():
    with open(version_file, 'r') as f:
        exec(compile(f.read(), version_file, 'exec'))
    return locals()['__version__']


# The full version, including alpha/beta/rc tags
release = get_version()

# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    'sphinx.ext.autodoc', 'sphinx.ext.napoleon', 'sphinx.ext.viewcode',
    'sphinx_markdown_tables', 'sphinx_copybutton', 'myst_parser'
]

autodoc_mock_imports = [
    'matplotlib', 'pycocotools', 'terminaltables', 'mmrazor.version',
    'mmcv.ops'
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# The suffix(es) of source filenames.
# You can specify multiple suffix as a list of string:
#
source_suffix = {
    '.rst': 'restructuredtext',
    '.md': 'markdown',
}

# The master toctree document.
master_doc = 'index'

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']

# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
# html_theme = 'sphinx_rtd_theme'
html_theme = 'pytorch_sphinx_theme'
html_theme_path = [pytorch_sphinx_theme.get_html_theme_path()]

html_theme_options = {
    'logo_url':
    'https://mmrazor.readthedocs.io/en/latest/',
    'menu': [
        {
            'name': 'GitHub',
            'url': 'https://github.com/open-mmlab/mmrazor'
        },
        {
            'name':
            'Upstream',
            'children': [{
                'name':
                'MMCV',
                'url':
                'https://github.com/open-mmlab/mmcv',
                'description':
                'Foundational library for computer vision'
            }, {
                'name':
                'MMDetection',
                'url':
                'https://github.com/open-mmlab/mmdetection',
                'description':
                'Object detection toolbox and benchmark'
            }, {
                'name':
                'MMClassification',
                'url':
                'https://github.com/open-mmlab/mmclassification',
                'description':
                'Image classification toolbox and benchmark'
            }, {
                'name':
                'MMSegmentation',
                'url':
                'https://github.com/open-mmlab/mmsegmentation',
                'description':
                'Semantic segmentation toolbox and benchmark'
            }]
        },
    ],
    'menu_lang':
    'en'
}

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']
html_css_files = ['css/readthedocs.css']

language = 'en'

# -- Extension configuration -------------------------------------------------
# Ignore >>> when copying code
copybutton_prompt_text = r'>>> |\.\.\. '
copybutton_prompt_is_regexp = True

# Enable ::: for my_st
myst_enable_extensions = ['colon_fence']
myst_heading_anchors = 4

docs/en/conf.py
Gumpest/AvatarKD
from pathlib import Path
from typing import Union


def gen_md_from_configs(config_root_dir: Union[Path, str],
                        target_md_path: Union[Path, str] = 'model_zoo.md',
                        prefix: str = '') -> None:

    def to_path(p: Union[Path, str]) -> Path:
        if isinstance(p, Path):
            return p
        if isinstance(p, str):
            return Path(p)
        raise ValueError(f'Unsupported type: {type(p)}')

    config_root_dir = to_path(config_root_dir)
    target_md_path = to_path(target_md_path)

    readme_path_list = []
    for readme_path in config_root_dir.rglob('README.md'):
        if readme_path.exists():
            config_name = readme_path.parent.name
            path = prefix / readme_path
            readme_path_list.append((config_name, path.as_posix()))

    with target_md_path.open('w', encoding='utf8') as f:
        f.write('# Model Zoo\n\n')
        f.write('## Baselines\n\n')

        for name, path in readme_path_list:
            f.write(f'### {name.upper()}\n\n')
            f.write(
                f'Please refer to [{name.upper()}]({path}) for details.\n\n')


if __name__ == '__main__':
    gen_md_from_configs('configs', 'docs/en/model_zoo.md', '/')

docs/en/gen_model_zoo.py
Gumpest/AvatarKD
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import os
import subprocess  # noqa: F401
import sys

import pytorch_sphinx_theme

sys.path.insert(0, os.path.abspath('../..'))

# -- Project information -----------------------------------------------------

project = 'MMRazor'
copyright = '2021-2030, OpenMMLab'
author = 'MMRazor Authors'
version_file = '../../mmrazor/version.py'


def get_version():
    with open(version_file, 'r') as f:
        exec(compile(f.read(), version_file, 'exec'))
    return locals()['__version__']


# The full version, including alpha/beta/rc tags
release = get_version()

# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    'sphinx.ext.autodoc', 'sphinx.ext.napoleon', 'sphinx.ext.viewcode',
    'sphinx_markdown_tables', 'sphinx_copybutton', 'myst_parser'
]

autodoc_mock_imports = [
    'matplotlib', 'pycocotools', 'terminaltables', 'mmrazor.version',
    'mmcv.ops'
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# The suffix(es) of source filenames.
# You can specify multiple suffix as a list of string:
#
source_suffix = {
    '.rst': 'restructuredtext',
    '.md': 'markdown',
}

# The master toctree document.
master_doc = 'index'

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']

# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
# html_theme = 'sphinx_rtd_theme'
html_theme = 'pytorch_sphinx_theme'
html_theme_path = [pytorch_sphinx_theme.get_html_theme_path()]

html_theme_options = {
    'logo_url':
    'https://mmrazor.readthedocs.io/zh_CN/latest/',
    'menu': [
        {
            'name': 'GitHub',
            'url': 'https://github.com/open-mmlab/mmrazor'
        },
        {
            'name':
            '上游库',
            'children': [{
                'name': 'MMCV',
                'url': 'https://github.com/open-mmlab/mmcv',
                'description': '基础视觉库'
            }, {
                'name': 'MMDetection',
                'url': 'https://github.com/open-mmlab/mmdetection',
                'description': '目标检测工具箱'
            }, {
                'name': 'MMClassification',
                'url': 'https://github.com/open-mmlab/mmclassification',
                'description': '图像分类工具箱'
            }, {
                'name': 'MMSegmentation',
                'url': 'https://github.com/open-mmlab/mmsegmentation',
                'description': '语义分割工具箱'
            }]
        },
    ],
    'menu_lang':
    'cn'
}

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']
html_css_files = ['css/readthedocs.css']

language = 'zh_CN'

# -- Extension configuration -------------------------------------------------
# Ignore >>> when copying code
copybutton_prompt_text = r'>>> |\.\.\. '
copybutton_prompt_is_regexp = True

# Enable ::: for my_st
myst_enable_extensions = ['colon_fence']
myst_heading_anchors = 4

docs/zh_cn/conf.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .mmcls import *  # noqa: F401,F403
from .mmdet import *  # noqa: F401,F403
from .mmseg import *  # noqa: F401,F403
from .utils import init_random_seed, set_random_seed  # noqa: F401

__all__ = ['init_random_seed', 'set_random_seed']

mmrazor/apis/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .inference import init_mmcls_model
from .train import set_random_seed, train_mmcls_model

__all__ = ['train_mmcls_model', 'init_mmcls_model', 'set_random_seed']

mmrazor/apis/mmcls/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import warnings
from typing import Dict, Optional, Union

import mmcv
from mmcv.runner import load_checkpoint
from torch import nn

from mmrazor.models import build_algorithm


def init_mmcls_model(config: Union[str, mmcv.Config],
                     checkpoint: Optional[str] = None,
                     device: str = 'cuda:0',
                     cfg_options: Optional[Dict] = None) -> nn.Module:
    """Initialize a mmcls model from config file.

    Args:
        config (str or :obj:`mmcv.Config`): Config file path or the config
            object.
        checkpoint (str, optional): Checkpoint path. If left as None, the model
            will not load any weights.
        cfg_options (dict): cfg_options to override some settings in the used
            config.

    Returns:
        nn.Module: The constructed classifier.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {type(config)}')
    if cfg_options is not None:
        config.merge_from_dict(cfg_options)

    model_cfg = config.algorithm.architecture.model
    model_cfg.pretrained = None
    algorithm = build_algorithm(config.algorithm)
    model = algorithm.architecture.model

    if checkpoint is not None:
        # Mapping the weights to GPU may cause unexpected video memory leak
        # which refers to https://github.com/open-mmlab/mmdetection/pull/6405
        checkpoint = load_checkpoint(algorithm, checkpoint, map_location='cpu')
        if 'CLASSES' in checkpoint.get('meta', {}):
            model.CLASSES = checkpoint['meta']['CLASSES']
        else:
            from mmcls.datasets import ImageNet
            warnings.simplefilter('once')
            warnings.warn('Class names are not saved in the checkpoint\'s '
                          'meta data, use imagenet by default.')
            model.CLASSES = ImageNet.CLASSES
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()

    return model

mmrazor/apis/mmcls/inference.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import random
import warnings

import numpy as np
import torch
from mmcls.core import DistOptimizerHook
from mmcls.datasets import build_dataloader, build_dataset
from mmcls.utils import get_root_logger
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import EpochBasedRunner, Fp16OptimizerHook, build_runner
from mmcv.runner.hooks import DistEvalHook, EvalHook

# Differences from mmclassification.
from mmrazor.core.distributed_wrapper import DistributedDataParallelWrapper
from mmrazor.core.hooks import DistSamplerSeedHook
from mmrazor.core.optimizer import build_optimizers
from mmrazor.datasets.utils import split_dataset
from mmrazor.utils import find_latest_checkpoint


def set_random_seed(seed, deterministic=False):
    """Import `set_random_seed` function here was deprecated in v0.3 and will
    be removed in v0.5.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set ``torch.backends.cudnn.deterministic``
            to True and ``torch.backends.cudnn.benchmark`` to False.
            Default: False.
    """
    warnings.warn(
        'Deprecated in v0.3 and will be removed in v0.5, '
        'please import `set_random_seed` directly from `mmrazor.apis`',
        category=DeprecationWarning)

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def train_mmcls_model(model,
                      dataset,
                      cfg,
                      distributed=False,
                      validate=False,
                      timestamp=None,
                      device='cuda',
                      meta=None):
    """Copy from mmclassification and modify some codes.

    This is an ugly implementation, and will be deprecated in the future. In
    the future, there will be only one train api and no longer distinguish
    between mmclassificaiton, mmsegmentation or mmdetection.
    """
    logger = get_root_logger()
    # Difference from mmclassification.
    # Split dataset.
    if cfg.data.pop('split', False):
        train_dataset = dataset[0]
        dataset[0] = split_dataset(train_dataset)

    loader_cfg = dict(
        # cfg.gpus will be ignored if distributed
        num_gpus=len(cfg.gpu_ids),
        dist=distributed,
        round_up=True,
        seed=cfg.get('seed'),
        sampler_cfg=cfg.get('sampler', None),
    )
    # The overall dataloader settings
    loader_cfg.update({
        k: v
        for k, v in cfg.data.items() if k not in [
            'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
            'test_dataloader'
        ]
    })
    # The specific dataloader settings
    train_loader_cfg = {**loader_cfg, **cfg.data.get('train_dataloader', {})}

    # Difference from mmclassification.
    # Build multi dataloaders according the splited datasets.
    data_loaders = list()
    for dset in dataset:
        if isinstance(dset, list):
            data_loader = [
                build_dataloader(item_ds, **train_loader_cfg)
                for item_ds in dset
            ]
        else:
            data_loader = build_dataloader(dset, **train_loader_cfg)

        data_loaders.append(data_loader)

    # put model on gpus
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        if cfg.get('use_ddp_wrapper', False):
            # Difference from mmclassification.
            # In some algorithms, the ``optimizer.step()`` is executed in
            # ``train_step``. To rebuilt reducer buckets rightly, there need to
            # use DistributedDataParallelWrapper.
            model = DistributedDataParallelWrapper(
                model,
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
        else:
            # Sets the ``find_unused_parameters`` parameter in
            # torch.nn.parallel.DistributedDataParallel
            model = MMDistributedDataParallel(
                model.cuda(),
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
    else:
        if device == 'cuda':
            model = MMDataParallel(
                model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
        elif device == 'cpu':
            warnings.warn(
                'The argument `device` is deprecated. To use cpu to train, '
                'please refers to https://mmclassification.readthedocs.io/en'
                '/latest/getting_started.html#train-a-model')
            model = model.cpu()
        else:
            raise ValueError(F'unsupported device name {device}.')

    # build optimizers
    # Difference from mmclassification.
    # In some algorithms, there will be multi optimizers.
    optimizer = build_optimizers(model, cfg.optimizer)

    if cfg.get('runner') is None:
        cfg.runner = {
            'type': 'EpochBasedRunner',
            'max_epochs': cfg.total_epochs
        }
        warnings.warn(
            'config is now expected to have a ``runner`` section, '
            'please set ``runner`` in your config.', UserWarning)

    runner = build_runner(
        cfg.runner,
        default_args=dict(
            model=model,
            batch_processor=None,
            optimizer=optimizer,
            work_dir=cfg.work_dir,
            logger=logger,
            meta=meta))

    # an ugly walkaround to make the .log and .log.json filenames the same
    runner.timestamp = timestamp

    # fp16 setting
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(
            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
    elif (distributed and cfg.optimizer_config is not None
          and 'type' not in cfg.optimizer_config):
        optimizer_config = DistOptimizerHook(**cfg.optimizer_config)
    else:
        optimizer_config = cfg.optimizer_config

    # register hooks
    runner.register_training_hooks(
        cfg.lr_config,
        optimizer_config,
        cfg.checkpoint_config,
        cfg.log_config,
        cfg.get('momentum_config', None),
        custom_hooks_config=cfg.get('custom_hooks', None))

    if distributed:
        if isinstance(runner, EpochBasedRunner):
            # Difference from mmclassification.
            # MMRazor's ``DistSamplerSeedHook`` could process multi dataloaders
            runner.register_hook(DistSamplerSeedHook())

    # register eval hooks
    if validate:
        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
        val_loader_cfg = {
            **loader_cfg,
            'shuffle': False,  # Not shuffle by default
            'sampler_cfg': None,  # Not use sampler by default
            **cfg.data.get('val_dataloader', {}),
        }
        val_dataloader = build_dataloader(val_dataset, **val_loader_cfg)
        eval_cfg = cfg.get('evaluation', {})

        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
        eval_hook = DistEvalHook if distributed else EvalHook
        # ``EvalHook`` needs to be executed after ``IterTimerHook``.
        # Otherwise, it will cause a bug if use ``IterBasedRunner``.
        # Refers to https://github.com/open-mmlab/mmcv/issues/1261
        runner.register_hook(
            eval_hook(val_dataloader, **eval_cfg), priority='LOW')

    resume_from = None
    if cfg.resume_from is None and cfg.get('auto_resume'):
        resume_from = find_latest_checkpoint(cfg.work_dir)
    if resume_from is not None:
        cfg.resume_from = resume_from

    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
    runner.run(data_loaders, cfg.workflow)

mmrazor/apis/mmcls/train.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmdet
except (ImportError, ModuleNotFoundError):
    mmdet = None

if mmdet is not None:
    from .inference import init_mmdet_model
    from .train import set_random_seed, train_mmdet_model

    __all__ = ['train_mmdet_model', 'init_mmdet_model', 'set_random_seed']

mmrazor/apis/mmdet/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import warnings
from typing import Dict, Optional, Union

import mmcv
from mmcv.runner import load_checkpoint
from mmdet.core import get_classes
from torch import nn

from mmrazor.models import build_algorithm


def init_mmdet_model(config: Union[str, mmcv.Config],
                     checkpoint: Optional[str] = None,
                     device: str = 'cuda:0',
                     cfg_options: Optional[Dict] = None) -> nn.Module:
    """Initialize a mmdet model from config file.

    Args:
        config (str or :obj:`mmcv.Config`): Config file path or the config
            object.
        checkpoint (str, optional): Checkpoint path. If left as None, the model
            will not load any weights.
        cfg_options (dict): Options to override some settings in the used
            config.

    Returns:
        nn.Module: The constructed detector.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {type(config)}')
    if cfg_options is not None:
        config.merge_from_dict(cfg_options)

    model_cfg = config.algorithm.architecture.model
    if 'pretrained' in model_cfg:
        model_cfg.pretrained = None
    elif 'init_cfg' in model_cfg.backbone:
        model_cfg.backbone.init_cfg = None

    model_cfg.train_cfg = None
    algorithm = build_algorithm(config.algorithm)
    model = algorithm.architecture.model

    if checkpoint is not None:
        checkpoint = load_checkpoint(algorithm, checkpoint, map_location='cpu')
        if 'CLASSES' in checkpoint.get('meta', {}):
            model.CLASSES = checkpoint['meta']['CLASSES']
        else:
            warnings.simplefilter('once')
            warnings.warn('Class names are not saved in the checkpoint\'s '
                          'meta data, use COCO classes by default.')
            model.CLASSES = get_classes('coco')
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()

    return model

mmrazor/apis/mmdet/inference.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import random
import warnings

import numpy as np
import torch
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (EpochBasedRunner, Fp16OptimizerHook, OptimizerHook,
                         build_runner)
from mmdet.core import DistEvalHook, EvalHook
from mmdet.datasets import (build_dataloader, build_dataset,
                            replace_ImageToTensor)
from mmdet.utils import get_root_logger

from mmrazor.core.distributed_wrapper import DistributedDataParallelWrapper
from mmrazor.core.hooks import DistSamplerSeedHook
from mmrazor.core.optimizer import build_optimizers
from mmrazor.utils import find_latest_checkpoint


def set_random_seed(seed, deterministic=False):
    """Import `set_random_seed` function here was deprecated in v0.3 and will
    be removed in v0.5.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set ``torch.backends.cudnn.deterministic``
            to True and ``torch.backends.cudnn.benchmark`` to False.
            Default: False.
    """
    warnings.warn(
        'Deprecated in v0.3 and will be removed in v0.5, '
        'please import `set_random_seed` directly from `mmrazor.apis`',
        category=DeprecationWarning)

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def train_mmdet_model(model,
                      dataset,
                      cfg,
                      distributed=False,
                      validate=False,
                      timestamp=None,
                      meta=None):
    """Copy from mmdetection and modify some codes.

    This is an ugly implementation, and will be deprecated in the future. In
    the future, there will be only one train api and no longer distinguish
    between mmclassificaiton, mmsegmentation or mmdetection.
    """

    logger = get_root_logger(cfg.log_level)

    # prepare data loaders
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]

    if 'imgs_per_gpu' in cfg.data:
        logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. '
                       'Please use "samples_per_gpu" instead')
        if 'samples_per_gpu' in cfg.data:
            logger.warning(
                f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and '
                f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"'
                f'={cfg.data.imgs_per_gpu} is used in this experiments')
        else:
            logger.warning(
                'Automatically set "samples_per_gpu"="imgs_per_gpu"='
                f'{cfg.data.imgs_per_gpu} in this experiments')
        cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu

    runner_type = 'EpochBasedRunner' if 'runner' not in cfg else cfg.runner[
        'type']
    data_loader = [
        build_dataloader(
            ds,
            cfg.data.samples_per_gpu,
            cfg.data.workers_per_gpu,
            # cfg.gpus will be ignored if distributed
            num_gpus=len(cfg.gpu_ids),
            dist=distributed,
            seed=cfg.seed,
            runner_type=runner_type,
            persistent_workers=cfg.data.get('persistent_workers', False))
        for ds in dataset
    ]

    # put model on gpus
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        if cfg.get('use_ddp_wrapper', False):
            # Difference from mmdetection.
            # In some algorithms, the ``optimizer.step()`` is executed in
            # ``train_step``. To rebuilt reducer buckets rightly, there need to
            # use DistributedDataParallelWrapper.
            model = DistributedDataParallelWrapper(
                model,
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
        else:
            # Sets the ``find_unused_parameters`` parameter in
            # torch.nn.parallel.DistributedDataParallel
            model = MMDistributedDataParallel(
                model.cuda(),
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
    else:
        model = MMDataParallel(model, device_ids=cfg.gpu_ids)

    # build optimizers
    # Difference from mmdetection.
    # In some algorithms, there will be multi optimizers.
    optimizer = build_optimizers(model, cfg.optimizer)

    # build runner
    if 'runner' not in cfg:
        cfg.runner = {
            'type': 'EpochBasedRunner',
            'max_epochs': cfg.total_epochs
        }
        warnings.warn(
            'config is now expected to have a ``runner`` section, '
            'please set ``runner`` in your config.', UserWarning)
    else:
        if 'total_epochs' in cfg:
            assert cfg.total_epochs == cfg.runner.max_epochs

    runner = build_runner(
        cfg.runner,
        default_args=dict(
            model=model,
            optimizer=optimizer,
            work_dir=cfg.work_dir,
            logger=logger,
            meta=meta))

    # an ugly workaround to make .log and .log.json filenames the same
    runner.timestamp = timestamp

    # fp16 setting
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(
            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
    elif distributed and cfg.optimizer_config is not None \
            and 'type' not in cfg.optimizer_config:
        optimizer_config = OptimizerHook(**cfg.optimizer_config)
    else:
        optimizer_config = cfg.optimizer_config

    # register hooks
    runner.register_training_hooks(
        cfg.lr_config,
        optimizer_config,
        cfg.checkpoint_config,
        cfg.log_config,
        cfg.get('momentum_config', None),
        custom_hooks_config=cfg.get('custom_hooks', None))

    if distributed:
        if isinstance(runner, EpochBasedRunner):
            # Difference from mmclassification.
            # MMRazor's ``DistSamplerSeedHook`` could process multi dataloaders
            runner.register_hook(DistSamplerSeedHook())

    # register eval hooks
    if validate:
        # Support batch_size > 1 in validation
        val_samples_per_gpu = cfg.data.val.pop('samples_per_gpu', 1)
        if val_samples_per_gpu > 1:
            # Replace 'ImageToTensor' to 'DefaultFormatBundle'
            cfg.data.val.pipeline = replace_ImageToTensor(
                cfg.data.val.pipeline)
        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
        val_dataloader = build_dataloader(
            val_dataset,
            samples_per_gpu=val_samples_per_gpu,
            workers_per_gpu=cfg.data.workers_per_gpu,
            dist=distributed,
            shuffle=False)
        eval_cfg = cfg.get('evaluation', {})
        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
        eval_hook = DistEvalHook if distributed else EvalHook
        runner.register_hook(
            eval_hook(val_dataloader, **eval_cfg), priority='LOW')

    resume_from = None
    if cfg.resume_from is None and cfg.get('auto_resume'):
        resume_from = find_latest_checkpoint(cfg.work_dir)
    if resume_from is not None:
        cfg.resume_from = resume_from

    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
    runner.run(data_loader, cfg.workflow)

mmrazor/apis/mmdet/train.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmseg
except (ImportError, ModuleNotFoundError):
    mmseg = None

if mmseg:
    from .inference import init_mmseg_model
    from .train import set_random_seed, train_mmseg_model

    __all__ = ['train_mmseg_model', 'init_mmseg_model', 'set_random_seed']

mmrazor/apis/mmseg/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Union

import mmcv
from mmcv.runner import load_checkpoint
from torch import nn

from mmrazor.models import build_algorithm


def init_mmseg_model(config: Union[str, mmcv.Config],
                     checkpoint: Optional[str] = None,
                     device: str = 'cuda:0') -> nn.Module:
    """Initialize a mmseg model from config file.

    Args:
        config (str or :obj:`mmcv.Config`): Config file path or the config
            object.
        checkpoint (str, optional): Checkpoint path. If left as None, the model
            will not load any weights.
        device (str, optional) CPU/CUDA device option. Default 'cuda:0'.
            Use 'cpu' for loading model on CPU.
    Returns:
        nn.Module: The constructed segmentor.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        'but got {}'.format(type(config)))

    model_cfg = config.algorithm.architecture.model
    model_cfg.pretrained = None
    model_cfg.train_cfg = None
    algorithm = build_algorithm(config.algorithm)
    model = algorithm.architecture.model

    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
        model.CLASSES = checkpoint['meta']['CLASSES']
        model.PALETTE = checkpoint['meta']['PALETTE']
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()

    return model

mmrazor/apis/mmseg/inference.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import random
import warnings

import numpy as np
import torch
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import EpochBasedRunner, build_runner
from mmseg.core import DistEvalHook, EvalHook
from mmseg.datasets import build_dataloader, build_dataset
from mmseg.utils import get_root_logger

from mmrazor.core.distributed_wrapper import DistributedDataParallelWrapper
from mmrazor.core.hooks import DistSamplerSeedHook
from mmrazor.core.optimizer import build_optimizers
from mmrazor.utils import find_latest_checkpoint


def set_random_seed(seed, deterministic=False):
    """Import `set_random_seed` function here was deprecated in v0.3 and will
    be removed in v0.5.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set ``torch.backends.cudnn.deterministic``
            to True and ``torch.backends.cudnn.benchmark`` to False.
            Default: False.
    """
    warnings.warn(
        'Deprecated in v0.3 and will be removed in v0.5, '
        'please import `set_random_seed` directly from `mmrazor.apis`',
        category=DeprecationWarning)

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def train_mmseg_model(model,
                      dataset,
                      cfg,
                      distributed=False,
                      validate=False,
                      timestamp=None,
                      meta=None):
    """Copy from mmsegmentation and modify some codes.

    This is an ugly implementation, and will be deprecated in the future. In
    the future, there will be only one train api and no longer distinguish
    between mmclassificaiton, mmsegmentation or mmdetection.
    """
    logger = get_root_logger(cfg.log_level)

    # prepare data loaders
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
    # The default loader config
    loader_cfg = dict(
        # cfg.gpus will be ignored if distributed
        num_gpus=len(cfg.gpu_ids),
        dist=distributed,
        seed=cfg.seed,
        drop_last=True)
    # The overall dataloader settings
    loader_cfg.update({
        k: v
        for k, v in cfg.data.items() if k not in [
            'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
            'test_dataloader'
        ]
    })

    # The specific dataloader settings
    train_loader_cfg = {**loader_cfg, **cfg.data.get('train_dataloader', {})}
    data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]

    # put model on gpus
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        if cfg.get('use_ddp_wrapper', False):
            # Difference from mmsegmentation.
            # In some algorithms, the ``optimizer.step()`` is executed in
            # ``train_step``. To rebuilt reducer buckets rightly, there need to
            # use DistributedDataParallelWrapper.
            model = DistributedDataParallelWrapper(
                model,
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
        else:
            # Sets the ``find_unused_parameters`` parameter in
            # torch.nn.parallel.DistributedDataParallel
            model = MMDistributedDataParallel(
                model.cuda(),
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
    else:
        model = MMDataParallel(model, device_ids=cfg.gpu_ids)

    # build optimizers
    # Difference from mmdetection.
    # In some algorithms, there will be multi optimizers.
    optimizer = build_optimizers(model, cfg.optimizer)

    # build runner
    if cfg.get('runner') is None:
        cfg.runner = {'type': 'IterBasedRunner', 'max_iters': cfg.total_iters}
        warnings.warn(
            'config is now expected to have a ``runner`` section, '
            'please set ``runner`` in your config.', UserWarning)

    runner = build_runner(
        cfg.runner,
        default_args=dict(
            model=model,
            batch_processor=None,
            optimizer=optimizer,
            work_dir=cfg.work_dir,
            logger=logger,
            meta=meta))

    # register hooks
    runner.register_training_hooks(
        cfg.lr_config,
        cfg.optimizer_config,
        cfg.checkpoint_config,
        cfg.log_config,
        cfg.get('momentum_config', None),
        custom_hooks_config=cfg.get('custom_hooks', None))

    # an ugly walkaround to make the .log and .log.json filenames the same
    runner.timestamp = timestamp

    if distributed:
        if isinstance(runner, EpochBasedRunner):
            runner.register_hook(DistSamplerSeedHook())

    # register eval hooks
    if validate:
        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
        # The specific dataloader settings
        val_loader_cfg = {
            **loader_cfg,
            'samples_per_gpu': 1,
            'shuffle': False,  # Not shuffle by default
            **cfg.data.get('val_dataloader', {}),
        }
        val_dataloader = build_dataloader(val_dataset, **val_loader_cfg)
        eval_cfg = cfg.get('evaluation', {})
        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
        eval_hook = DistEvalHook if distributed else EvalHook
        # In this PR (https://github.com/open-mmlab/mmcv/pull/1193), the
        # priority of IterTimerHook has been modified from 'NORMAL' to 'LOW'.
        runner.register_hook(
            eval_hook(val_dataloader, **eval_cfg), priority='LOW')

    resume_from = None
    if cfg.resume_from is None and cfg.get('auto_resume'):
        resume_from = find_latest_checkpoint(cfg.work_dir)
    if resume_from is not None:
        cfg.resume_from = resume_from

    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
    runner.run(data_loaders, cfg.workflow)

mmrazor/apis/mmseg/train.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import random

import numpy as np
import torch
import torch.distributed as dist
from mmcv.runner import get_dist_info


def init_random_seed(seed=None, device='cuda'):
    """Initialize random seed.

    If the seed is not set, the seed will be automatically randomized,
    and then broadcast to all processes to prevent some potential bugs.
    Args:
        seed (int, Optional): The seed. Default to None.
        device (str): The device where the seed will be put on.
            Default to 'cuda'.
    Returns:
        int: Seed to be used.
    """
    if seed is not None:
        return seed

    # Make sure all ranks share the same random seed to prevent
    # some potential bugs. Please refer to
    # https://github.com/open-mmlab/mmdetection/issues/6339
    rank, world_size = get_dist_info()
    seed = np.random.randint(2**31)
    if world_size == 1:
        return seed

    if rank == 0:
        random_num = torch.tensor(seed, dtype=torch.int32, device=device)
    else:
        random_num = torch.tensor(0, dtype=torch.int32, device=device)
    dist.broadcast(random_num, src=0)
    return random_num.item()


def set_random_seed(seed: int, deterministic: bool = False) -> None:
    """Set random seed.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set ``torch.backends.cudnn.deterministic``
            to True and ``torch.backends.cudnn.benchmark`` to False.
            Default: False.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

mmrazor/apis/utils.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .builder import build_searcher
from .distributed_wrapper import DistributedDataParallelWrapper
from .hooks import *  # noqa: F401,F403
from .optimizer import *  # noqa: F401,F403
from .runners import *  # noqa: F401,F403
from .searcher import *  # noqa: F401,F403
from .utils import *  # noqa: F401,F403

__all__ = ['DistributedDataParallelWrapper', 'build_searcher']

mmrazor/core/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.utils import Registry, build_from_cfg

SEARCHERS = Registry('search')


def build_searcher(cfg, default_args=None):
    """Build searcher."""
    return build_from_cfg(cfg, SEARCHERS, default_args=default_args)

mmrazor/core/builder.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from mmcv.parallel import MODULE_WRAPPERS as MMCV_MODULE_WRAPPERS
from mmcv.parallel import MMDistributedDataParallel
from mmcv.parallel.scatter_gather import scatter_kwargs
from mmcv.utils import Registry
from torch.cuda._utils import _get_device_index

MODULE_WRAPPERS = Registry('module wrapper', parent=MMCV_MODULE_WRAPPERS)


@MODULE_WRAPPERS.register_module()
class DistributedDataParallelWrapper(nn.Module):
    """A DistributedDataParallel wrapper for models in MMediting.

    In MMedting, there is a need to wrap different modules in the models
    with separate DistributedDataParallel. Otherwise, it will cause
    errors for GAN training.
    More specific, the GAN model, usually has two sub-modules:
    generator and discriminator. If we wrap both of them in one
    standard DistributedDataParallel, it will cause errors during training,
    because when we update the parameters of the generator (or discriminator),
    the parameters of the discriminator (or generator) is not updated, which is
    not allowed for DistributedDataParallel.
    So we design this wrapper to separately wrap DistributedDataParallel
    for generator and discriminator.
    In this wrapper, we perform two operations:
    1. Wrap the modules in the models with separate MMDistributedDataParallel.
        Note that only modules with parameters will be wrapped.
    2. Do scatter operation for 'forward', 'train_step' and 'val_step'.
    Note that the arguments of this wrapper is the same as those in
    ``torch.nn.parallel.distributed.DistributedDataParallel``.

    Args:
        module (nn.Module): Module that needs to be wrapped.
        device_ids (list[int | ``torch.device``]): Same as that in
            ``torch.nn.parallel.distributed.DistributedDataParallel``.
        dim (int, optional): Same as that in the official scatter function in
            pytorch. Defaults to 0.
        broadcast_buffers (bool): Same as that in
            ``torch.nn.parallel.distributed.DistributedDataParallel``.
            Defaults to False.
        find_unused_parameters (bool, optional): Same as that in
            ``torch.nn.parallel.distributed.DistributedDataParallel``.
            Traverse the autograd graph of all tensors contained in returned
            value of the wrapped module’s forward function. Defaults to False.
        kwargs (dict): Other arguments used in
            ``torch.nn.parallel.distributed.DistributedDataParallel``.
    """

    def __init__(self,
                 module,
                 device_ids,
                 dim=0,
                 broadcast_buffers=False,
                 find_unused_parameters=False,
                 **kwargs):
        super().__init__()
        assert len(device_ids) == 1, (
            'Currently, DistributedDataParallelWrapper only supports one'
            'single CUDA device for each process.'
            f'The length of device_ids must be 1, but got {len(device_ids)}.')
        self.module = module
        self.dim = dim
        self.to_ddp(
            device_ids=device_ids,
            dim=dim,
            broadcast_buffers=broadcast_buffers,
            find_unused_parameters=find_unused_parameters,
            **kwargs)
        self.output_device = _get_device_index(device_ids[0], True)

    def to_ddp(self, device_ids, dim, broadcast_buffers,
               find_unused_parameters, **kwargs):
        """Wrap models with separate MMDistributedDataParallel.

        It only wraps the modules with parameters.
        """
        for name, module in self.module._modules.items():
            if next(module.parameters(), None) is None:
                module = module.cuda()
            elif all(not p.requires_grad for p in module.parameters()):
                module = module.cuda()
            else:
                module = MMDistributedDataParallel(
                    module.cuda(),
                    device_ids=device_ids,
                    dim=dim,
                    broadcast_buffers=broadcast_buffers,
                    find_unused_parameters=find_unused_parameters,
                    **kwargs)
            self.module._modules[name] = module

    def scatter(self, inputs, kwargs, device_ids):
        """Scatter function.

        Args:
            inputs (Tensor): Input Tensor.
            kwargs (dict): Args for
                ``mmcv.parallel.scatter_gather.scatter_kwargs``.
            device_ids (int): Device id.
        """
        return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)

    def forward(self, *inputs, **kwargs):
        """Forward function.

        Args:
            inputs (tuple): Input data.
            kwargs (dict): Args for
                ``mmcv.parallel.scatter_gather.scatter_kwargs``.
        """
        inputs, kwargs = self.scatter(inputs, kwargs,
                                      [torch.cuda.current_device()])
        return self.module(*inputs[0], **kwargs[0])

    def train_step(self, *inputs, **kwargs):
        """Train step function.

        Args:
            inputs (Tensor): Input Tensor.
            kwargs (dict): Args for
                ``mmcv.parallel.scatter_gather.scatter_kwargs``.
        """
        inputs, kwargs = self.scatter(inputs, kwargs,
                                      [torch.cuda.current_device()])
        output = self.module.train_step(*inputs[0], **kwargs[0])
        return output

    def val_step(self, *inputs, **kwargs):
        """Validation step function.

        Args:
            inputs (tuple): Input data.
            kwargs (dict): Args for ``scatter_kwargs``.
        """
        inputs, kwargs = self.scatter(inputs, kwargs,
                                      [torch.cuda.current_device()])
        output = self.module.val_step(*inputs[0], **kwargs[0])
        return output

mmrazor/core/distributed_wrapper.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .drop_path_prob import DropPathProbHook
from .sampler_seed import DistSamplerSeedHook
from .search_subnet import SearchSubnetHook

__all__ = ['DistSamplerSeedHook', 'DropPathProbHook', 'SearchSubnetHook']

mmrazor/core/hooks/__init__.py
Gumpest/AvatarKD
# Copyright (c) Open-MMLab. All rights reserved.

from mmcv.cnn.bricks import DropPath
from mmcv.runner import HOOKS, Hook


@HOOKS.register_module()
class DropPathProbHook(Hook):
    """Set drop_path_prob periodically.

    Args:
        max_prob (float): The max probability of dropping.
        interval (int): The saving period. If ``by_epoch=True``, interval
            indicates epochs, otherwise it indicates iterations.
            Default: -1, which means "never".
        by_epoch (bool): Saving checkpoints by epoch or by iteration.
            Default: True.
    """

    def __init__(self, max_prob, interval=-1, by_epoch=True, **kwargs):
        self.interval = interval
        self.by_epoch = by_epoch
        self.max_prob = max_prob
        assert self.by_epoch

    def before_train_epoch(self, runner):
        """Executed in before_train_epoch stage."""
        if not self.by_epoch:
            return

        if self.every_n_epochs(
                runner, self.interval) or (self.save_last
                                           and self.is_last_epoch(runner)):
            cur_epoch = runner.epoch
            max_epoch = runner._max_epochs
            drop_prob = self.max_prob * (cur_epoch * 1.0) / (max_epoch * 1.0)
            for module in runner.model.modules():
                if isinstance(module, DropPath):
                    module.drop_prob = drop_prob
            runner.logger.info(f'Set drop_prob to {drop_prob} \
                    at {runner.epoch + 1} epochs')

mmrazor/core/hooks/drop_path_prob.py
Gumpest/AvatarKD
# Copyright (c) Open-MMLab. All rights reserved.
from mmcv.runner import Hook


# @HOOKS.register_module()
class DistSamplerSeedHook(Hook):
    """Data-loading sampler for distributed training.

    When distributed training, it is only useful in conjunction with
    :obj:`EpochBasedRunner`, while :obj:``IterBasedRunner`` achieves the same
    purpose with :obj:`IterLoader`.
    """

    def before_epoch(self, runner):
        """Executed in before_epoch stage."""
        if hasattr(runner.data_loader, '_dataloaders'):
            data_loaders = runner.data_loader._dataloaders
        else:
            data_loaders = [runner.data_loader]

        for data_loader in data_loaders:

            if hasattr(data_loader.sampler, 'set_epoch'):
                # in case the data loader uses ``SequentialSampler`` in Pytorch
                data_loader.sampler.set_epoch(runner.epoch)
            elif hasattr(data_loader.batch_sampler.sampler, 'set_epoch'):
                # batch sampler in pytorch warps the sampler as its attributes.
                data_loader.batch_sampler.sampler.set_epoch(runner.epoch)

mmrazor/core/hooks/sampler_seed.py
Gumpest/AvatarKD
# Copyright (c) Open-MMLab. All rights reserved.
import os

from mmcv.runner import HOOKS, Hook, master_only


@HOOKS.register_module()
class SearchSubnetHook(Hook):
    """Save checkpoints periodically.

    Args:
        interval (int): The saving period. If ``by_epoch=True``, interval
            indicates epochs, otherwise it indicates iterations.
            Default: -1, which means "never".
        by_epoch (bool): Saving checkpoints by epoch or by iteration.
            Default: True.
        out_dir (str, optional): The directory to save checkpoints. If not
            specified, ``runner.work_dir`` will be used by default.
        max_keep_ckpts (int, optional): The maximum checkpoints to keep.
            In some cases we want only the latest few checkpoints and would
            like to delete old ones to save the disk space.
            Default: -1, which means unlimited.
        save_last (bool): Whether to force the last checkpoint to be saved
            regardless of interval.
    """

    def __init__(self,
                 interval=-1,
                 by_epoch=True,
                 out_dir=None,
                 max_keep_ckpts=-1,
                 save_last=True,
                 **kwargs):
        self.interval = interval
        self.by_epoch = by_epoch

        self.out_dir = out_dir
        self.max_keep_ckpts = max_keep_ckpts
        self.save_last = save_last
        self.args = kwargs

    def before_run(self, runner):
        """Executed in before_run stage."""
        if not self.out_dir:
            self.out_dir = runner.work_dir

    def after_train_epoch(self, runner):
        """Executed in after_train_epoch stage."""
        if not self.by_epoch:
            return

        # save checkpoint for following cases:
        # 1. every ``self.interval`` epochs
        # 2. reach the last epoch of training
        if self.every_n_epochs(
                runner, self.interval) or (self.save_last
                                           and self.is_last_epoch(runner)):
            runner.logger.info(f'Saving subnet at {runner.epoch + 1} epochs')
            self._search_subnet(runner)

    @master_only
    def _search_subnet(self, runner):
        """Save the current checkpoint and delete unwanted checkpoint."""
        runner.search_subnet(self.out_dir, **self.args)
        if runner.meta is not None:
            if self.by_epoch:
                cur_subnet_filename = self.args.get(
                    'filename_tmpl', 'epoch_{}.yaml').format(runner.epoch + 1)
            else:
                cur_subnet_filename = self.args.get(
                    'filename_tmpl', 'iter_{}.yaml').format(runner.iter + 1)
            runner.meta.setdefault('hook_msgs', dict())
            runner.meta['hook_msgs']['last_subnet'] = os.path.join(
                self.out_dir, cur_subnet_filename)
        # remove other checkpoints
        if self.max_keep_ckpts > 0:
            if self.by_epoch:
                name = 'epoch_{}.yaml'
                current_subnet = runner.epoch + 1
            else:
                name = 'iter_{}.yaml'
                current_subnet = runner.iter + 1
            redundant_subnets = range(
                current_subnet - self.max_keep_subnets * self.interval, 0,
                -self.interval)
            filename_tmpl = self.args.get('filename_tmpl', name)
            for _step in redundant_subnets:
                subnet_path = os.path.join(self.out_dir,
                                           filename_tmpl.format(_step))
                if os.path.exists(subnet_path):
                    os.remove(subnet_path)
                else:
                    break

    def after_train_iter(self, runner):
        """Executed in after_train_iter stage."""
        if self.by_epoch:
            return

        # save checkpoint for following cases:
        # 1. every ``self.interval`` iterations
        # 2. reach the last iteration of training
        if self.every_n_iters(
                runner, self.interval) or (self.save_last
                                           and self.is_last_iter(runner)):
            runner.logger.info(
                f'Saving subnet at {runner.iter + 1} iterations')

            self._search_subnet(runner)

mmrazor/core/hooks/search_subnet.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.runner import build_optimizer


def build_optimizers(model, cfgs):
    """Build multiple optimizers from configs. If `cfgs` contains several dicts
    for optimizers, then a dict for each constructed optimizers will be
    returned. If `cfgs` only contains one optimizer config, the constructed
    optimizer itself will be returned. For example, 1) Multiple optimizer
    configs:
    code-block::

        optimizer_cfg = dict(
            model1=dict(type='SGD', lr=lr),
            model2=dict(type='SGD', lr=lr))

    The return dict is
    ``dict('model1': torch.optim.Optimizer, 'model2': torch.optim.Optimizer)``
    2) Single optimizer config:
    .. code-block::

        optimizer_cfg = dict(type='SGD', lr=lr)

    The return is ``torch.optim.Optimizer``.
    Args:
        model (:obj:`nn.Module`): The model with parameters to be optimized.
        cfgs (dict): The config dict of the optimizer.
    Returns:
        dict[:obj:`torch.optim.Optimizer`] | :obj:`torch.optim.Optimizer`:
            The initialized optimizers.
    """
    optimizers = {}
    if hasattr(model, 'module'):
        model = model.module
    # determine whether 'cfgs' has several dicts for optimizers
    is_dict_of_dict = True
    for key, cfg in cfgs.items():
        if not isinstance(cfg, dict):
            is_dict_of_dict = False

    if is_dict_of_dict:
        for key, cfg in cfgs.items():
            cfg_ = cfg.copy()
            module = getattr(model, key)
            optimizers[key] = build_optimizer(module, cfg_)
        return optimizers

    return build_optimizer(model, cfgs)

mmrazor/core/optimizer/builder.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .epoch_based_runner import MultiLoaderEpochBasedRunner
from .iter_based_runner import MultiLoaderIterBasedRunner

__all__ = ['MultiLoaderEpochBasedRunner', 'MultiLoaderIterBasedRunner']

mmrazor/core/runners/__init__.py
Gumpest/AvatarKD
# Copyright (c) Open-MMLab. All rights reserved.
import os.path as osp
import platform
import shutil
import time
from functools import partial

import mmcv
import yaml
from mmcv.runner import HOOKS, RUNNERS, EpochBasedRunner

from ..utils import set_lr


class EpochMultiLoader:
    """Multi loaders based on epoch."""

    def __init__(self, dataloaders):
        self._dataloaders = dataloaders
        self.iter_loaders = [iter(loader) for loader in self._dataloaders]

    @property
    def num_loaders(self):
        """The number of dataloaders."""
        return len(self._dataloaders)

    def __iter__(self):
        """Return self when executing __iter__."""
        return self

    def __next__(self):
        """Get next iter's data."""
        data = tuple([next(loader) for loader in self.iter_loaders])

        return data

    def __len__(self):
        """Get the length of loader."""
        return min([len(loader) for loader in self._dataloaders])


@RUNNERS.register_module()
class MultiLoaderEpochBasedRunner(EpochBasedRunner):
    """Multi Dataloaders EpochBasedRunner.

    There are three differences from EpochBaseRunner： 1）Support load data from
    multi dataloaders. 2) Support freeze some optimizer's lr update when runner
    has multi optimizers. 3) Add ``search_subnet`` api.
    """

    def train(self, data_loader, **kwargs):
        """Rewrite the ``train`` of ``EpochBasedRunner``."""
        self.model.train()
        self.mode = 'train'

        # the only difference from EpochBasedRunner's ``train``
        if isinstance(data_loader, list):
            self.data_loader = EpochMultiLoader(data_loader)
        else:
            self.data_loader = data_loader

        self._max_iters = self._max_epochs * len(self.data_loader)
        self.call_hook('before_train_epoch')
        time.sleep(2)  # Prevent possible deadlock during epoch transition
        for i, data_batch in enumerate(self.data_loader):
            self._inner_iter = i
            self.call_hook('before_train_iter')
            self.run_iter(data_batch, train_mode=True, **kwargs)
            self.call_hook('after_train_iter')
            self._iter += 1

        self.call_hook('after_train_epoch')

        self._epoch += 1

    def register_lr_hook(self, lr_config):
        """Resister a hook for setting learning rate.

        Args:
            lr_config (dict): Config for setting learning rate.
        """
        if lr_config is None:
            return
        elif isinstance(lr_config, dict):
            assert 'policy' in lr_config
            policy_type = lr_config.pop('policy')
            if 'freeze_optimizers' in lr_config:
                freeze_optimizers = lr_config.pop('freeze_optimizers')
            else:
                freeze_optimizers = []
            # If the type of policy is all in lower case, e.g., 'cyclic',
            # then its first letter will be capitalized, e.g., to be 'Cyclic'.
            # This is for the convenient usage of Lr updater.
            # Since this is not applicable for ``CosineAnnealingLrUpdater``,
            # the string will not be changed if it contains capital letters.
            if policy_type == policy_type.lower():
                policy_type = policy_type.title()
            hook_type = policy_type + 'LrUpdaterHook'
            lr_config['type'] = hook_type
            hook = mmcv.build_from_cfg(lr_config, HOOKS)

            # modify the hook's ``_set_lr``
            # the only difference from BasedRunner's ``register_lr_hook``
            hook._set_lr = partial(set_lr, freeze_optimizers=freeze_optimizers)
        else:
            hook = lr_config

        self.register_hook(hook, priority=10)

    def search_subnet(self,
                      out_dir,
                      filename_tmpl='epoch_{}.yaml',
                      create_symlink=True):
        """Search the best subnet.

        Args:
            out_dir (str): The directory that subnets are saved.
            filename_tmpl (str, optional): The subnet filename template,
                which contains a placeholder for the epoch number.
                Defaults to 'epoch_{}.yaml'.
            create_symlink (bool, optional): Whether to create a symlink
                "latest.yaml" to point to the latest subnet.
                Defaults to True.
        """

        filename = filename_tmpl.format(self.epoch + 1)
        filepath = osp.join(out_dir, filename)

        algorithm = self.model.module if hasattr(self.model,
                                                 'module') else self.model
        mutator = algorithm.mutator.module if hasattr(
            algorithm.mutator, 'module') else algorithm.mutator
        subnet = mutator.search_subnet()
        with open(filepath, 'w') as f:
            yaml.dump(subnet, f)
        # in some environments, ``os.symlink`` is not supported, you may need
        # to set ``create_symlink`` to False
        if create_symlink:
            dst_file = osp.join(out_dir, 'latest.yaml')
            if platform.system() != 'Windows':
                mmcv.symlink(filename, dst_file)
            else:
                shutil.copy(filepath, dst_file)

mmrazor/core/runners/epoch_based_runner.py
Gumpest/AvatarKD
# Copyright (c) Open-MMLab. All rights reserved.
import os.path as osp
import platform
import shutil
import time
import warnings
from functools import partial

import mmcv
import yaml
from mmcv.runner import HOOKS, RUNNERS, IterBasedRunner
from mmcv.runner.utils import get_host_info

from ..utils import set_lr


class IterMultiLoader:
    """Multi loaders based on iter."""

    def __init__(self, dataloaders):
        self._dataloaders = dataloaders if isinstance(dataloaders,
                                                      list) else [dataloaders]
        self.iter_loaders = [iter(loader) for loader in self._dataloaders]
        self._epoch = 0

    @property
    def epoch(self):
        """The property of the class."""
        return self._epoch

    @property
    def num_loaders(self):
        """The number of dataloaders."""
        return len(self._dataloaders)

    def __next__(self):
        """Get next iter's data."""
        try:
            data = tuple([next(loader) for loader in self.iter_loaders])
        except StopIteration:
            self._epoch += 1
            for loader in self._dataloaders:
                if hasattr(loader.sampler, 'set_epoch'):
                    loader.sampler.set_epoch(self._epoch)
            self.iter_loader = [iter(loader) for loader in self._dataloaders]
            data = tuple([next(loader) for loader in self.iter_loaders])

        return data

    def __len__(self):
        """Get the length of loader."""
        return min([len(loader) for loader in self._dataloaders])


@RUNNERS.register_module()
class MultiLoaderIterBasedRunner(IterBasedRunner):
    """Multi Dataloaders IterBasedRunner.

    There are three differences from IterBasedRunner 1）Support load data from
    multi dataloaders. 2) Support freeze some optimizer's lr update when runner
    has multi optimizers. 3) Add ``search_subnet`` api.
    """

    def run(self, data_loaders, workflow, max_iters=None, **kwargs):
        """Start running.

        Args:
            data_loaders (list[:obj:`DataLoader`]): Dataloaders for training
                and validation.
            workflow (list[tuple]): A list of (phase, iters) to specify the
                running order and iterations. E.g, [('train', 10000),
                ('val', 1000)] means running 10000 iterations for training and
                1000 iterations for validation, iteratively.
            max_iters (int): Specify the max iters.
        """
        assert isinstance(data_loaders, list)
        assert mmcv.is_list_of(workflow, tuple)
        assert len(data_loaders) == len(workflow)
        if max_iters is not None:
            warnings.warn(
                'setting max_iters in run is deprecated, '
                'please set max_iters in runner_config', DeprecationWarning)
            self._max_iters = max_iters
        assert self._max_iters is not None, (
            'max_iters must be specified during instantiation')

        work_dir = self.work_dir if self.work_dir is not None else 'NONE'
        self.logger.info('Start running, host: %s, work_dir: %s',
                         get_host_info(), work_dir)
        self.logger.info('Hooks will be executed in the following order:\n%s',
                         self.get_hook_info())
        self.logger.info('workflow: %s, max: %d iters', workflow,
                         self._max_iters)
        self.call_hook('before_run')

        # the only difference from IterBasedRunner's ``train``
        iter_loaders = [IterMultiLoader(x) for x in data_loaders]

        self.call_hook('before_epoch')

        while self.iter < self._max_iters:
            for i, flow in enumerate(workflow):
                self._inner_iter = 0
                mode, iters = flow
                if not isinstance(mode, str) or not hasattr(self, mode):
                    raise ValueError(
                        'runner has no method named "{}" to run a workflow'.
                        format(mode))
                iter_runner = getattr(self, mode)
                for _ in range(iters):
                    if mode == 'train' and self.iter >= self._max_iters:
                        break
                    iter_runner(iter_loaders[i], **kwargs)

        time.sleep(1)  # wait for some hooks like loggers to finish
        self.call_hook('after_epoch')
        self.call_hook('after_run')

    def register_lr_hook(self, lr_config):
        """Resister a hook for setting learning rate.

        Args:
            lr_config (dict): Config for setting learning rate.
        """
        if lr_config is None:
            return
        elif isinstance(lr_config, dict):
            assert 'policy' in lr_config
            policy_type = lr_config.pop('policy')
            if 'freeze_optimizers' in lr_config:
                freeze_optimizers = lr_config.pop('freeze_optimizers')
            else:
                freeze_optimizers = []
            # If the type of policy is all in lower case, e.g., 'cyclic',
            # then its first letter will be capitalized, e.g., to be 'Cyclic'.
            # This is for the convenient usage of Lr updater.
            # Since this is not applicable for ``CosineAnnealingLrUpdater``,
            # the string will not be changed if it contains capital letters.
            if policy_type == policy_type.lower():
                policy_type = policy_type.title()
            hook_type = policy_type + 'LrUpdaterHook'
            lr_config['type'] = hook_type
            hook = mmcv.build_from_cfg(lr_config, HOOKS)

            # modify the hook's ``_set_lr``
            # the only difference from BasedRunner's ``register_lr_hook``
            hook._set_lr = partial(set_lr, freeze_optimizers=freeze_optimizers)
        else:
            hook = lr_config

        self.register_hook(hook, priority=10)

    def search_subnet(self,
                      out_dir,
                      filename_tmpl='epoch_{}.yaml',
                      create_symlink=True):
        """Search the best subnet.

        Args:
            out_dir (str): The directory that subnets are saved.
            filename_tmpl (str, optional): The subnet filename template,
                which contains a placeholder for the epoch number.
                Defaults to 'epoch_{}.yaml'.
            create_symlink (bool, optional): Whether to create a symlink
                "latest.yaml" to point to the latest subnet.
                Defaults to True.
        """

        filename = filename_tmpl.format(self.epoch + 1)
        filepath = osp.join(out_dir, filename)

        # save_checkpoint(self.model, filepath, optimizer=optimizer, meta=meta)
        subnet = self.model.module.search_subnet()
        with open(filepath, 'w') as f:
            yaml.dump(subnet, f)
        # in some environments, ``os.symlink`` is not supported, you may need
        # to set ``create_symlink`` to False
        if create_symlink:
            dst_file = osp.join(out_dir, 'latest.yaml')
            if platform.system() != 'Windows':
                mmcv.symlink(filename, dst_file)
            else:
                shutil.copy(filepath, dst_file)

mmrazor/core/runners/iter_based_runner.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .evolution_search import EvolutionSearcher
from .greedy_search import GreedySearcher

__all__ = ['EvolutionSearcher', 'GreedySearcher']

mmrazor/core/searcher/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import random
import time

import mmcv.fileio
from mmcv.runner import get_dist_info

from ..builder import SEARCHERS
from ..utils import broadcast_object_list


@SEARCHERS.register_module()
class EvolutionSearcher():
    """Implement of evolution search.

    Args:
        algorithm (:obj:`torch.nn.Module`): Algorithm to be used.
        dataloader (nn.Dataloader): Pytorch data loader.
        test_fn (function): Test api to used for evaluation.
        work_dir (str): Working direction is to save search result and log.
        logger (logging.Logger): To log info in search stage.
        candidate_pool_size (int): The length of candidate pool.
        candidate_top_k (int): Specify top k candidates based on scores.
        constraints (dict): Constraints to be used for screening candidates.
        metrics (str): Metrics to be used for evaluating candidates.
        metric_options (str): Options to be used for metrics.
        score_key (str): To be used for specifying one metric from evaluation
            results.
        max_epoch (int): Specify max epoch to end evolution search.
        num_mutation (int): The number of candidates got by mutation.
        num_crossover (int): The number of candidates got by crossover.
        mutate_prob (float): The probability of mutation.
        resume_from (str): Specify the path of saved .pkl file for resuming
            searching
    """

    def __init__(self,
                 algorithm,
                 dataloader,
                 test_fn,
                 work_dir,
                 logger,
                 candidate_pool_size=50,
                 candidate_top_k=10,
                 constraints=dict(flops=330 * 1e6),
                 metrics=None,
                 metric_options=None,
                 score_key='accuracy_top-1',
                 max_epoch=20,
                 num_mutation=25,
                 num_crossover=25,
                 mutate_prob=0.1,
                 resume_from=None,
                 **search_kwargs):

        if not hasattr(algorithm, 'module'):
            raise NotImplementedError('Do not support searching with cpu.')
        if num_mutation + num_crossover > candidate_pool_size:
            raise ValueError(
                f'The sum of num_mutation({num_mutation}) and '
                f'num_crossover({num_crossover}) should not be '
                f'greater than candidate_pool_size({candidate_pool_size}).')
        self.algorithm = algorithm.module
        self.algorithm_for_test = algorithm
        self.dataloader = dataloader
        self.constraints = constraints
        self.metrics = metrics
        self.metric_options = metric_options
        self.score_key = score_key
        self.candidate_pool = list()
        self.candidate_pool_size = candidate_pool_size
        self.max_epoch = max_epoch
        self.test_fn = test_fn
        self.candidate_top_k = candidate_top_k
        self.num_mutation = num_mutation
        self.num_crossover = num_crossover
        self.mutate_prob = mutate_prob
        self.top_k_candidates_with_score = dict()
        self.candidate_pool_with_score = dict()
        self.work_dir = work_dir
        self.resume_from = resume_from
        self.logger = logger

    def check_constraints(self):
        """Check whether is beyond constraints.

        Returns:
            bool: The result of checking.
        """
        flops = self.algorithm.get_subnet_flops()
        if flops < self.constraints['flops']:
            return True
        else:
            return False

    def update_top_k(self):
        """Update top k candidates."""
        self.top_k_candidates_with_score.update(self.candidate_pool_with_score)
        self.top_k_candidates_with_score = dict(
            sorted(
                self.top_k_candidates_with_score.items(),
                key=lambda x: x[0],
                reverse=True))
        keys = list(self.top_k_candidates_with_score.keys())
        new_dict = dict()
        for k in keys[:self.candidate_top_k]:
            new_dict[k] = self.top_k_candidates_with_score[k]
        self.top_k_candidates_with_score = new_dict.copy()

    def search(self):
        """Execute the pipeline of evolution search."""
        epoch_start = 0
        if self.resume_from is not None:
            searcher_resume = mmcv.fileio.load(self.resume_from)
            for k in searcher_resume.keys():
                setattr(self, k, searcher_resume[k])
            epoch_start = int(searcher_resume['epoch'])
            self.logger.info('#' * 100)
            self.logger.info(f'Resume from epoch: {epoch_start}')
            self.logger.info('#' * 100)
        self.logger.info('Experiment setting:')
        self.logger.info(f'candidate_pool_size: {self.candidate_pool_size}')
        self.logger.info(f'candidate_top_k: {self.candidate_top_k}')
        self.logger.info(f'num_crossover: {self.num_crossover}')
        self.logger.info(f'num_mutation: {self.num_mutation}')
        self.logger.info(f'mutate_prob: {self.mutate_prob}')
        self.logger.info(f'max_epoch: {self.max_epoch}')
        self.logger.info(f'score_key: {self.score_key}')
        self.logger.info(f'constraints: {self.constraints}')
        self.logger.info('#' * 100)

        rank = get_dist_info()[0]
        for epoch in range(epoch_start, self.max_epoch):
            if rank == 0:
                while len(self.candidate_pool) < self.candidate_pool_size:
                    candidate = \
                        self.algorithm.mutator.sample_subnet(searching=True)
                    self.algorithm.mutator.set_subnet(candidate)

                    if self.check_constraints():
                        self.candidate_pool.append(candidate)
            else:
                self.candidate_pool = [None] * self.candidate_pool_size
            broadcast_object_list(self.candidate_pool)

            for i, candidate in enumerate(self.candidate_pool):
                self.algorithm.mutator.set_subnet(candidate)
                outputs = self.test_fn(self.algorithm_for_test,
                                       self.dataloader)

                if rank == 0:
                    eval_result = self.dataloader.dataset.evaluate(
                        outputs, self.metrics, self.metric_options)
                    score = eval_result[self.score_key]
                    self.candidate_pool_with_score[score] = candidate
                    self.logger.info(f'Epoch:[{epoch + 1}/{self.max_epoch}] '
                                     f'Candidate:[{i + 1}/'
                                     f'{self.candidate_pool_size}] '
                                     f'Score:{score}')
            if rank == 0:
                scores_before = list(self.top_k_candidates_with_score.keys())
                self.logger.info(f'top k scores before update: '
                                 f'{scores_before}')
                self.update_top_k()
                scores_after = list(self.top_k_candidates_with_score.keys())
                self.logger.info(f'top k scores before update: '
                                 f'{scores_after}')

                mutation_candidates = list()
                max_mutate_iters = self.num_mutation * 10
                mutate_iter = 0
                while len(mutation_candidates) < self.num_mutation:
                    mutate_iter += 1
                    if mutate_iter > max_mutate_iters:
                        break
                    candidate = random.choice(
                        list(self.top_k_candidates_with_score.values()))
                    mutation = self.algorithm.mutator.mutation(
                        candidate, self.mutate_prob)
                    self.algorithm.mutator.set_subnet(mutation)
                    if self.check_constraints():
                        mutation_candidates.append(mutation)

                crossover_candidates = list()
                crossover_iter = 0
                max_crossover_iters = self.num_crossover * 10
                while len(crossover_candidates) < self.num_crossover:
                    crossover_iter += 1
                    if crossover_iter > max_crossover_iters:
                        break

                    random_candidate1 = random.choice(
                        list(self.top_k_candidates_with_score.values()))
                    random_candidate2 = random.choice(
                        list(self.top_k_candidates_with_score.values()))

                    crossover_candidate = \
                        self.algorithm.mutator.crossover(
                            random_candidate1, random_candidate2)
                    self.algorithm.mutator.set_subnet(crossover_candidate)
                    if self.check_constraints():
                        crossover_candidates.append(crossover_candidate)

                self.candidate_pool = (
                    mutation_candidates + crossover_candidates)

                save_for_resume = dict()
                save_for_resume['epoch'] = epoch + 1
                for k in ['candidate_pool', 'top_k_candidates_with_score']:
                    save_for_resume[k] = getattr(self, k)
                mmcv.fileio.dump(
                    save_for_resume,
                    osp.join(self.work_dir, f'search_epoch_{epoch + 1}.pkl'))
                self.logger.info(
                    f'Epoch:[{epoch + 1}/{self.max_epoch}], top1_score: '
                    f'{list(self.top_k_candidates_with_score.keys())[0]}')

        if rank == 0:
            final_subnet_dict = list(
                self.top_k_candidates_with_score.values())[0]
            self.algorithm.mutator.set_chosen_subnet(final_subnet_dict)
            final_subnet_dict_to_save = dict()
            for k in final_subnet_dict.keys():
                final_subnet_dict_to_save[k] = dict({
                    'chosen':
                    self.algorithm.mutator.search_spaces[k]['chosen']
                })
            timestamp_subnet = time.strftime('%Y%m%d_%H%M', time.localtime())
            save_name = f'final_subnet_{timestamp_subnet}.yaml'
            mmcv.fileio.dump(final_subnet_dict_to_save,
                             osp.join(self.work_dir, save_name))
            self.logger.info('Search finished and '
                             f'{save_name} saved in {self.work_dir}.')

mmrazor/core/searcher/evolution_search.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import os
import os.path as osp

import mmcv
import torch
from mmcv.runner import get_dist_info

from ..builder import SEARCHERS
from ..utils import broadcast_object_list


@SEARCHERS.register_module()
class GreedySearcher():
    """Search with the greedy algorithm.

    We start with the largest model and compare the network accuracy among
    the architectures where each layer is slimmed by one channel bin. We then
    greedily slim the layer with minimal performance drop. During the iterative
    slimming, we obtain optimized channel configurations under different
    resource constraints. We stop until reaching the strictest constraint
    (e.g., 200M FLOPs).

    Args:
        algorithm (:obj:`torch.nn.Module`): Specific implemented algorithm
         based specific task in mmRazor, eg: AutoSlim.
        dataloader (:obj:`torch.nn.Dataloader`): Pytorch data loader.
        target_flops (list): The target flops of the searched models.
        test_fn (callable): test a model with samples from a dataloader,
            and return the test results.
        work_dir (str): Output result file.
        logger (logging.Logger): To log info in search stage.
        max_channel_bins (int): The maximum number of channel bins in each
            layer. Note that each layer is slimmed by one channel bin.
        min_channel_bins (int): The minimum number of channel bins in each
            layer. Default to 1.
        metrics (str | list[str]): Metrics to be evaluated.
            Default value is ``accuracy``
        metric_options (dict, optional): Options for calculating metrics.
            Allowed keys are 'topk', 'thrs' and 'average_mode'.
            Defaults to None.
        score_key (str): The metric to judge the performance of a model.
            Defaults to `accuracy_top-1`.
        resume_from (str, optional): Specify the path of saved .pkl file for
            resuming searching. Defaults to None.
    """

    def __init__(self,
                 algorithm,
                 dataloader,
                 target_flops,
                 test_fn,
                 work_dir,
                 logger,
                 max_channel_bins,
                 min_channel_bins=1,
                 metrics='accuracy',
                 metric_options=None,
                 score_key='accuracy_top-1',
                 resume_from=None,
                 **search_kwargs):
        super(GreedySearcher, self).__init__()
        if not hasattr(algorithm, 'module'):
            raise NotImplementedError('Do not support searching with cpu.')

        self.algorithm = algorithm.module
        self.algorithm_for_test = algorithm
        self.dataloader = dataloader
        self.target_flops = sorted(target_flops, reverse=True)
        self.test_fn = test_fn
        self.work_dir = work_dir
        self.logger = logger
        self.max_channel_bins = max_channel_bins
        self.min_channel_bins = min_channel_bins
        self.metrics = metrics
        self.metric_options = metric_options
        self.score_key = score_key
        self.resume_from = resume_from

    def search(self):
        """Greedy Slimming."""
        algorithm = self.algorithm

        rank, _ = get_dist_info()

        if self.resume_from is not None:
            searcher_resume = mmcv.fileio.load(self.resume_from)  # a dict
            result_subnet = searcher_resume['result_subnet']
            result_flops = searcher_resume['result_flops']
            subnet = searcher_resume['subnet']
            flops = searcher_resume['flops']
            self.logger.info(f'Resume from subnet: {subnet}')
        else:
            result_subnet, result_flops = [], []
            # We start with the largest model
            algorithm.pruner.set_max_channel()
            max_subnet = algorithm.pruner.get_max_channel_bins(
                self.max_channel_bins)
            # channel_cfg
            subnet = max_subnet
            flops = algorithm.get_subnet_flops()

        for target in self.target_flops:
            if self.resume_from is not None and flops <= target:
                continue

            if flops <= target:
                algorithm.pruner.set_channel_bins(subnet,
                                                  self.max_channel_bins)
                channel_cfg = algorithm.pruner.export_subnet()
                result_subnet.append(channel_cfg)
                result_flops.append(flops)
                self.logger.info(f'Find model flops {flops} <= {target}')
                continue

            while flops > target:
                # search which layer needs to shrink
                best_score = None
                best_subnet = None

                # During distributed training, the order of ``subnet.keys()``
                # on different ranks may be different. So we need to sort it
                # first.
                for i, name in enumerate(sorted(subnet.keys())):
                    new_subnet = copy.deepcopy(subnet)
                    # we prune the very last channel bin
                    last_bin_ind = torch.where(new_subnet[name] == 1)[0][-1]
                    # The ``new_subnet`` on different ranks are the same,
                    # so we do not need to broadcast here.
                    new_subnet[name][last_bin_ind] = 0
                    if torch.sum(new_subnet[name]) < self.min_channel_bins:
                        # subnet is invalid
                        continue

                    algorithm.pruner.set_channel_bins(new_subnet,
                                                      self.max_channel_bins)

                    outputs = self.test_fn(self.algorithm_for_test,
                                           self.dataloader)
                    broadcast_scores = [None]
                    if rank == 0:
                        eval_result = self.dataloader.dataset.evaluate(
                            outputs, self.metrics, self.metric_options)
                        broadcast_scores = [eval_result[self.score_key]]

                    # Broadcasts scores in broadcast_scores to the whole
                    # group.
                    broadcast_object_list(broadcast_scores)
                    score = broadcast_scores[0]
                    self.logger.info(
                        f'Slimming group {name}, {self.score_key}: {score}')
                    if best_score is None or score > best_score:
                        best_score = score
                        best_subnet = new_subnet

                if best_subnet is None:
                    raise RuntimeError(
                        'Cannot find any valid model, check your '
                        'configurations.')

                subnet = best_subnet
                algorithm.pruner.set_channel_bins(subnet,
                                                  self.max_channel_bins)
                flops = algorithm.get_subnet_flops()
                self.logger.info(
                    f'Greedy find model, score: {best_score}, FLOPS: {flops}')

                save_for_resume = dict()
                save_for_resume['result_subnet'] = result_subnet
                save_for_resume['result_flops'] = result_flops
                save_for_resume['subnet'] = subnet
                save_for_resume['flops'] = flops
                mmcv.fileio.dump(save_for_resume,
                                 osp.join(self.work_dir, 'latest.pkl'))

            algorithm.pruner.set_channel_bins(subnet, self.max_channel_bins)
            channel_cfg = algorithm.pruner.export_subnet()
            result_subnet.append(channel_cfg)
            result_flops.append(flops)
            self.logger.info(f'Find model flops {flops} <= {target}')

        self.logger.info('Search models done.')

        if rank == 0:
            for flops, subnet in zip(result_flops, result_subnet):
                mmcv.fileio.dump(
                    subnet,
                    os.path.join(self.work_dir,
                                 'subnet_{}.yaml'.format(flops)))
            self.logger.info(f'Save searched results to {self.work_dir}')

mmrazor/core/searcher/greedy_search.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .broadcast import broadcast_object_list
from .lr import set_lr
from .utils import get_backend, get_default_group, get_rank, get_world_size

__all__ = [
    'broadcast_object_list', 'set_lr', 'get_world_size', 'get_rank',
    'get_backend', 'get_default_group'
]

mmrazor/core/utils/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import pickle
import warnings
from typing import Any, List, Optional, Tuple

import torch
from mmcv.utils import TORCH_VERSION, digit_version
from torch import Tensor
from torch import distributed as dist

from .utils import get_backend, get_default_group, get_rank, get_world_size


def _object_to_tensor(obj: Any) -> Tuple[Tensor, Tensor]:
    """Serialize picklable python object to tensor."""
    byte_storage = torch.ByteStorage.from_buffer(pickle.dumps(obj))
    # Do not replace `torch.ByteTensor` or `torch.LongTensor` with torch.tensor
    # and specifying dtype. Otherwise, it will cause 100X slowdown.
    # See: https://github.com/pytorch/pytorch/issues/65696
    byte_tensor = torch.ByteTensor(byte_storage)
    local_size = torch.LongTensor([byte_tensor.numel()])
    return byte_tensor, local_size


def _tensor_to_object(tensor: Tensor, tensor_size: int) -> Any:
    """Deserialize tensor to picklable python object."""
    buf = tensor.cpu().numpy().tobytes()[:tensor_size]
    return pickle.loads(buf)


def _broadcast_object_list(object_list: List[Any],
                           src: int = 0,
                           group: Optional[dist.ProcessGroup] = None) -> None:
    """Broadcast picklable objects in ``object_list`` to the whole group.

    Similar to :func:`broadcast`, but Python objects can be passed in. Note
    that all objects in ``object_list`` must be picklable in order to be
    broadcasted.
    """
    if dist.distributed_c10d._rank_not_in_group(group):
        return

    my_rank = get_rank()
    # Serialize object_list elements to tensors on src rank.
    if my_rank == src:
        tensor_list, size_list = zip(
            *[_object_to_tensor(obj) for obj in object_list])
        object_sizes_tensor = torch.cat(size_list)
    else:
        object_sizes_tensor = torch.empty(len(object_list), dtype=torch.long)

    # Current device selection.
    # To preserve backwards compatibility, ``device`` is ``None`` by default.
    # in which case we run current logic of device selection, i.e.
    # ``current_device`` is CUDA if backend is NCCL otherwise CPU device. In
    # the case it is not ``None`` we move the size and object tensors to be
    # broadcasted to this device.
    group_backend = get_backend(group)
    is_nccl_backend = group_backend == dist.Backend.NCCL
    current_device = torch.device('cpu')
    if is_nccl_backend:
        # See note about using torch.cuda.current_device() here in
        # docstring. We cannot simply use my_rank since rank == device is
        # not necessarily true.
        current_device = torch.device('cuda', torch.cuda.current_device())
        object_sizes_tensor = object_sizes_tensor.to(current_device)

    # Broadcast object sizes
    dist.broadcast(object_sizes_tensor, src=src, group=group)

    # Concatenate and broadcast serialized object tensors
    if my_rank == src:
        object_tensor = torch.cat(tensor_list)
    else:
        object_tensor = torch.empty(
            torch.sum(object_sizes_tensor).int().item(),
            dtype=torch.uint8,
        )

    if is_nccl_backend:
        object_tensor = object_tensor.to(current_device)
    dist.broadcast(object_tensor, src=src, group=group)
    # Deserialize objects using their stored sizes.
    offset = 0
    if my_rank != src:
        for i, obj_size in enumerate(object_sizes_tensor):
            obj_view = object_tensor[offset:offset + obj_size]
            obj_view = obj_view.type(torch.uint8)
            if obj_view.device != torch.device('cpu'):
                obj_view = obj_view.cpu()
            offset += obj_size
            object_list[i] = _tensor_to_object(obj_view, obj_size)


def broadcast_object_list(data: List[Any],
                          src: int = 0,
                          group: Optional[dist.ProcessGroup] = None) -> None:
    """Broadcasts picklable objects in ``object_list`` to the whole group.
    Similar to :func:`broadcast`, but Python objects can be passed in. Note
    that all objects in ``object_list`` must be picklable in order to be
    broadcasted.
    Note:
        Calling ``broadcast_object_list`` in non-distributed environment does
        nothing.
    Args:
        data (List[Any]): List of input objects to broadcast.
            Each object must be picklable. Only objects on the ``src`` rank
            will be broadcast, but each rank must provide lists of equal sizes.
        src (int): Source rank from which to broadcast ``object_list``.
        group: (ProcessGroup, optional): The process group to work on. If None,
            the default process group will be used. Default is ``None``.
        device (``torch.device``, optional): If not None, the objects are
            serialized and converted to tensors which are moved to the
            ``device`` before broadcasting. Default is ``None``.
    Note:
        For NCCL-based process groups, internal tensor representations of
        objects must be moved to the GPU device before communication starts.
        In this case, the used device is given by
        ``torch.cuda.current_device()`` and it is the user's responsibility to
        ensure that this is correctly set so that each rank has an individual
        GPU, via ``torch.cuda.set_device()``.
    Examples:
        >>> import torch
        >>> import mmrazor.core.utils as dist
        >>> # non-distributed environment
        >>> data = ['foo', 12, {1: 2}]
        >>> dist.broadcast_object_list(data)
        >>> data
        ['foo', 12, {1: 2}]
        >>> # distributed environment
        >>> # We have 2 process groups, 2 ranks.
        >>> if dist.get_rank() == 0:
        >>>     # Assumes world_size of 3.
        >>>     data = ["foo", 12, {1: 2}]  # any picklable object
        >>> else:
        >>>     data = [None, None, None]
        >>> dist.broadcast_object_list(data)
        >>> data
        ["foo", 12, {1: 2}]  # Rank 0
        ["foo", 12, {1: 2}]  # Rank 1
    """
    warnings.warn(
        '`broadcast_object_list` is now without return value, '
        'and it\'s input parameters are: `data`,`src` and '
        '`group`, but its function is similar to the old\'s', UserWarning)
    assert isinstance(data, list)

    if get_world_size(group) > 1:
        if group is None:
            group = get_default_group()

        if digit_version(TORCH_VERSION) >= digit_version('1.8.0'):
            dist.broadcast_object_list(data, src, group)
        else:
            _broadcast_object_list(data, src, group)

mmrazor/core/utils/broadcast.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
def set_lr(runner, lr_groups, freeze_optimizers=[]):
    """Set specified learning rate in optimizer."""
    if isinstance(runner.optimizer, dict):
        for k, optim in runner.optimizer.items():
            if k in freeze_optimizers:
                continue
            for param_group, lr in zip(optim.param_groups, lr_groups[k]):
                param_group['lr'] = lr
    else:
        for param_group, lr in zip(runner.optimizer.param_groups, lr_groups):
            param_group['lr'] = lr

mmrazor/core/utils/lr.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional

from torch import distributed as dist


def is_distributed() -> bool:
    """Return True if distributed environment has been initialized."""
    return dist.is_available() and dist.is_initialized()


def get_default_group() -> Optional[dist.ProcessGroup]:
    """Return default process group."""

    return dist.distributed_c10d._get_default_group()


def get_rank(group: Optional[dist.ProcessGroup] = None) -> int:
    """Return the rank of the given process group.

    Rank is a unique identifier assigned to each process within a distributed
    process group. They are always consecutive integers ranging from 0 to
    ``world_size``.
    Note:
        Calling ``get_rank`` in non-distributed environment will return 0.
    Args:
        group (ProcessGroup, optional): The process group to work on. If None,
            the default process group will be used. Defaults to None.
    Returns:
        int: Return the rank of the process group if in distributed
        environment, otherwise 0.
    """

    if is_distributed():
        # handle low versions of torch like 1.5.0 which does not support
        # passing in None for group argument
        if group is None:
            group = get_default_group()
        return dist.get_rank(group)
    else:
        return 0


def get_backend(group: Optional[dist.ProcessGroup] = None) -> Optional[str]:
    """Return the backend of the given process group.

    Note:
        Calling ``get_backend`` in non-distributed environment will return
        None.
    Args:
        group (ProcessGroup, optional): The process group to work on. The
            default is the general main process group. If another specific
            group is specified, the calling process must be part of
            :attr:`group`. Defaults to None.
    Returns:
        str or None: Return the backend of the given process group as a lower
        case string if in distributed environment, otherwise None.
    """
    if is_distributed():
        # handle low versions of torch like 1.5.0 which does not support
        # passing in None for group argument
        if group is None:
            group = get_default_group()
        return dist.get_backend(group)
    else:
        return None


def get_world_size(group: Optional[dist.ProcessGroup] = None) -> int:
    """Return the number of the given process group.

    Note:
        Calling ``get_world_size`` in non-distributed environment will return
        1.
    Args:
        group (ProcessGroup, optional): The process group to work on. If None,
            the default process group will be used. Defaults to None.
    Returns:
        int: Return the number of processes of the given process group if in
        distributed environment, otherwise 1.
    """
    if is_distributed():
        # handle low versions of torch like 1.5.0 which does not support
        # passing in None for group argument
        if group is None:
            group = get_default_group()
        return dist.get_world_size(group)
    else:
        return 1

mmrazor/core/utils/utils.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .utils import split_dataset

__all__ = ['split_dataset']

mmrazor/datasets/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from torch.utils.data import random_split


def split_dataset(dataset):
    dset_length = len(dataset)

    first_dset_length = dset_length // 2
    second_dset_length = dset_length - first_dset_length
    split_tuple = (first_dset_length, second_dset_length)
    first_dset, second_dset = random_split(dataset, split_tuple)

    first_dset.CLASSES = dataset.CLASSES
    second_dset.CLASSES = dataset.CLASSES

    return [first_dset, second_dset]

mmrazor/datasets/utils.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmdet
except (ImportError, ModuleNotFoundError):
    mmdet = None

try:
    import mmseg
except (ImportError, ModuleNotFoundError):
    mmseg = None

from mmcls.models import *  # noqa: F401,F403

from .algorithms import *  # noqa: F401,F403
from .architectures import *  # noqa: F401,F403
from .builder import (ALGORITHMS, ARCHITECTURES, DISTILLERS, LOSSES, MUTABLES,
                      MUTATORS, OPS, build_algorithm, build_architecture,
                      build_distiller, build_loss, build_mutable,
                      build_mutator, build_op)
from .distillers import *  # noqa: F401,F403
from .losses import *  # noqa: F401,F403
from .mutables import *  # noqa: F401,F403
from .mutators import *  # noqa: F401,F403
from .ops import *  # noqa: F401,F403
from .pruners import *  # noqa: F401,F403

if mmdet is not None:
    from mmdet.models import *  # noqa: F401,F403
if mmseg is not None:
    from mmseg.models import *  # noqa: F401,F403

__all__ = [
    'ALGORITHMS', 'MUTABLES', 'ARCHITECTURES', 'DISTILLERS', 'MUTATORS',
    'LOSSES', 'OPS', 'build_architecture', 'build_mutable', 'build_op',
    'build_mutator', 'build_algorithm', 'build_distiller', 'build_loss'
]

mmrazor/models/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .align_method_kd import AlignMethodDistill
from .autoslim import AutoSlim
from .darts import Darts
from .detnas import DetNAS
from .general_distill import GeneralDistill
from .spos import SPOS

__all__ = [
    'AutoSlim', 'AlignMethodDistill', 'Darts', 'SPOS', 'DetNAS',
    'GeneralDistill'
]

mmrazor/models/algorithms/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmrazor.models.builder import ALGORITHMS
from .general_distill import GeneralDistill


@ALGORITHMS.register_module()
class AlignMethodDistill(GeneralDistill):

    def __init__(self, **kwargs):
        super(AlignMethodDistill, self).__init__(**kwargs)

    def train_step(self, data, optimizer):

        with self.distiller.context_manager:
            outputs = super().train_step(data, optimizer)
        return outputs

mmrazor/models/algorithms/align_method_kd.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy

import mmcv
import torch
import torch.nn as nn
from mmcv.cnn import get_model_complexity_info
from torch.nn.modules.batchnorm import _BatchNorm

from mmrazor.models.builder import ALGORITHMS, build_pruner
from mmrazor.models.utils import add_prefix
from .base import BaseAlgorithm


@ALGORITHMS.register_module()
class AutoSlim(BaseAlgorithm):
    """AutoSlim: A one-shot architecture search for channel numbers.

    Please refer to the `paper <https://arxiv.org/abs/1903.11728>` for
    details.

    Args:
        num_sample_training (int): In each iteration we train the model at
            smallest width, largest width and (`num_sample_training` − 2)
            random widths. It should be no less than 2. Defaults to 4
        input_shape (tuple): Input shape used for calculation the flops of the
            supernet.
        bn_training_mode (bool): Whether set bn to training mode when model is
            set to eval mode. Note that in slimmable networks, accumulating
            different numbers of channels results in different feature means
            and variances, which further leads to inaccurate statistics of
            shared BN. Set ``bn_training_mode`` to True to use the feature
            means and variances in a batch.
    """

    def __init__(self,
                 num_sample_training=4,
                 input_shape=(3, 224, 224),
                 bn_training_mode=False,
                 **kwargs):

        super(AutoSlim, self).__init__(**kwargs)
        assert num_sample_training >= 2, \
            'num_sample_training should be no less than 2'
        self.num_sample_training = num_sample_training

        # set bn to training mode when model is set to eval mode
        self.bn_training_mode = bn_training_mode

        if input_shape is not None:
            self.input_shape = input_shape
            self._init_flops()

    def _init_pruner(self, pruner):
        """Build registered pruners and make preparations.

        Args:
            pruner (dict): The registered pruner to be used
                in the algorithm.
        """
        if pruner is None:
            self.pruner = None
            return

        # judge whether our StructurePruner can prune the architecture
        try:
            pseudo_pruner = build_pruner(pruner)
            pseudo_architecture = copy.deepcopy(self.architecture)
            pseudo_pruner.prepare_from_supernet(pseudo_architecture)
            subnet_dict = pseudo_pruner.sample_subnet()
            pseudo_pruner.set_subnet(subnet_dict)
            subnet_dict = pseudo_pruner.export_subnet()

            pseudo_pruner.deploy_subnet(pseudo_architecture, subnet_dict)
            pseudo_img = torch.randn(1, 3, 224, 224)
            pseudo_architecture.forward_dummy(pseudo_img)
        except RuntimeError:
            raise NotImplementedError('Our current StructurePruner does not '
                                      'support pruning this architecture. '
                                      'StructurePruner is not perfect enough '
                                      'to handle all the corner cases. We will'
                                      ' appreciate it if you create a issue.')

        self.pruner = build_pruner(pruner)

        if self.retraining:
            if isinstance(self.channel_cfg, dict):
                self.pruner.deploy_subnet(self.architecture, self.channel_cfg)
                self.deployed = True
            elif isinstance(self.channel_cfg, (list, tuple)):

                self.pruner.convert_switchable_bn(self.architecture,
                                                  len(self.channel_cfg))
                self.pruner.prepare_from_supernet(self.architecture)
            else:
                raise NotImplementedError
        else:
            self.pruner.prepare_from_supernet(self.architecture)

    def _init_flops(self):
        """Get flops information of the supernet."""
        flops_model = copy.deepcopy(self.architecture)
        flops_model.eval()
        if hasattr(flops_model, 'forward_dummy'):
            flops_model.forward = flops_model.forward_dummy
        else:
            raise NotImplementedError(
                'FLOPs counter is currently not currently supported with {}'.
                format(flops_model.__class__.__name__))

        flops, params = get_model_complexity_info(
            flops_model, self.input_shape, print_per_layer_stat=False)
        flops_lookup = dict()
        for name, module in flops_model.named_modules():
            flops = getattr(module, '__flops__', 0)
            flops_lookup[name] = flops
        del (flops_model)

        for name, module in self.architecture.named_modules():
            module.__flops__ = flops_lookup[name]

    def get_subnet_flops(self):
        """A hacky way to get flops information of a subnet."""
        flops = 0
        last_out_mask_ratio = None
        for name, module in self.architecture.named_modules():
            if type(module) in [
                    nn.Conv2d, mmcv.cnn.bricks.Conv2d, nn.Linear,
                    mmcv.cnn.bricks.Linear
            ]:
                in_mask_ratio = float(module.in_mask.sum() /
                                      module.in_mask.numel())
                out_mask_ratio = float(module.out_mask.sum() /
                                       module.out_mask.numel())
                flops += module.__flops__ * in_mask_ratio * out_mask_ratio
                last_out_mask_ratio = out_mask_ratio
            elif type(module) == nn.BatchNorm2d:
                out_mask_ratio = float(module.out_mask.sum() /
                                       module.out_mask.numel())
                flops += module.__flops__ * out_mask_ratio
                last_out_mask_ratio = out_mask_ratio
            elif type(module) in [
                    nn.ReLU, nn.PReLU, nn.ELU, nn.LeakyReLU, nn.ReLU6
            ]:

                assert last_out_mask_ratio, 'An activate module can not be ' \
                                            'the first module of a network.'
                flops += module.__flops__ * last_out_mask_ratio

        return round(flops)

    def train_step(self, data, optimizer):
        """Train step function.

        This function implements the standard training iteration for
        autoslim pretraining and retraining.

        Args:
            data (dict): Input data from dataloader.
            optimizer (:obj:`torch.optim.Optimizer`): The optimizer to
                accumulate gradient
        """
        optimizer.zero_grad()

        losses = dict()
        if not self.retraining:
            assert self.pruner is not None

            self.pruner.set_max_channel()
            if self.distiller is not None:
                max_model_losses = self.distiller.exec_teacher_forward(
                    self.architecture, data)
            else:
                max_model_losses = self(**data)
            losses.update(add_prefix(max_model_losses, 'max_model'))
            max_model_loss, _ = self._parse_losses(max_model_losses)
            max_model_loss.backward()

            self.pruner.set_min_channel()
            if self.distiller is not None:
                self.distiller.exec_student_forward(self.architecture, data)
                min_model_losses = self.distiller.compute_distill_loss(data)
            else:
                min_model_losses = self(**data)
            losses.update(add_prefix(min_model_losses, 'min_model'))
            min_model_loss, _ = self._parse_losses(min_model_losses)
            min_model_loss.backward()

            for i in range(self.num_sample_training - 2):
                subnet_dict = self.pruner.sample_subnet()
                self.pruner.set_subnet(subnet_dict)
                if self.distiller is not None:
                    self.distiller.exec_student_forward(
                        self.architecture, data)
                    model_losses = self.distiller.compute_distill_loss(data)
                    losses.update(
                        add_prefix(model_losses,
                                   'prune_model{}_distiller'.format(i + 1)))
                else:
                    model_losses = self(**data)
                    losses.update(
                        add_prefix(model_losses,
                                   'prune_model{}'.format(i + 1)))
                model_loss, _ = self._parse_losses(model_losses)
                model_loss.backward()
        else:
            if self.deployed:
                # Only one subnet retrains. The supernet has already deploy
                model_losses = self(**data)
                losses.update(add_prefix(model_losses, 'prune_model'))
                model_loss, _ = self._parse_losses(model_losses)
                model_loss.backward()
            else:
                # More than one subnet retraining together
                assert isinstance(self.channel_cfg, (list, tuple))
                for i, subnet in enumerate(self.channel_cfg):
                    self.pruner.switch_subnet(subnet, i)
                    model_losses = self(**data)
                    losses.update(
                        add_prefix(model_losses,
                                   'prune_model_{}'.format(i + 1)))
                    model_loss, _ = self._parse_losses(model_losses)
                    model_loss.backward()

        # TODO: clip grad norm
        optimizer.step()

        loss, log_vars = self._parse_losses(losses)
        outputs = dict(
            loss=loss, log_vars=log_vars, num_samples=len(data['img'].data))

        return outputs

    def train(self, mode=True):
        """Overwrite the train method in ``nn.Module`` to set ``nn.BatchNorm``
        to training mode when model is set to eval mode when
        ``self.bn_training_mode`` is ``True``.

        Args:
            mode (bool): whether to set training mode (``True``) or evaluation
                mode (``False``). Default: ``True``.
        """
        super(AutoSlim, self).train(mode)
        if not mode and self.bn_training_mode:
            for module in self.modules():
                if isinstance(module, _BatchNorm):
                    module.training = True

mmrazor/models/algorithms/autoslim.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from collections import OrderedDict

import mmcv.fileio
import torch
import torch.distributed as dist
from mmcv.runner import BaseModule

from mmrazor.models.builder import (ALGORITHMS, build_architecture,
                                    build_distiller, build_mutator,
                                    build_pruner)


@ALGORITHMS.register_module()
class BaseAlgorithm(BaseModule):
    """Base class for algorithms, it consists of two main parts: architecture
    and algorithm components.

    Args:
        architecture (dict): Config for architecture to be slimmed.
        mutator (dict): Config for mutator, which is an algorithm component
            for NAS.
        pruner (dict): Config for pruner, which is an algorithm component
            for pruning.
        distiller (dict): Config for pruner, which is an algorithm component
            for knowledge distillation.
        retraining (bool): Whether is in retraining stage, if False,
            it is in pre-training stage.
        init_cfg (dict): Init config for ``BaseModule``.
        mutable_cfg (dict): Config for mutable of the subnet searched out,
            it will be needed in retraining stage.
        channel_cfg (dict): Config for channel of the subnet searched out,
            it will be needed in retraining stage.
    """

    def __init__(
        self,
        architecture,
        mutator=None,
        pruner=None,
        distiller=None,
        retraining=False,
        init_cfg=None,
        mutable_cfg=None,
        channel_cfg=None,
    ):
        super(BaseAlgorithm, self).__init__(init_cfg)
        self.retraining = retraining
        if self.retraining:
            self.mutable_cfg = self.load_subnet(mutable_cfg)
            self.channel_cfg = self.load_subnet(channel_cfg)
        self.architecture = build_architecture(architecture)

        self.deployed = False
        self._init_mutator(mutator)
        self._init_pruner(pruner)
        self._init_distiller(distiller)

    def load_subnet(self, subnet_path):
        """Load subnet searched out in search stage.

        Args:
            subnet_path (str | list[str] | tuple(str)): The path of saved
                subnet file, its suffix should be .yaml.
            There may be several subnet searched out in some algorithms.

        Returns:
            dict | list[dict]: Config(s) for subnet(s) searched out.
        """
        if subnet_path is None:
            return

        if isinstance(subnet_path, str):
            cfg = mmcv.fileio.load(subnet_path)
        elif isinstance(subnet_path, (list, tuple)):
            cfg = list()
            for path in subnet_path:
                cfg.append(mmcv.fileio.load(path))
        else:
            raise NotImplementedError

        return cfg

    def _init_mutator(self, mutator):
        """Build registered mutators and make preparations.

        Args:
            mutator (dict): Config for mutator, which is an algorithm component
                for NAS.
        """
        if mutator is None:
            self.mutator = None
            return
        self.mutator = build_mutator(mutator)
        self.mutator.prepare_from_supernet(self.architecture)
        if self.retraining:
            if isinstance(self.mutable_cfg, dict):
                self.mutator.deploy_subnet(self.architecture, self.mutable_cfg)
                self.deployed = True
            else:
                raise NotImplementedError

    def _init_pruner(self, pruner):
        """Build registered pruners and make preparations.

        Args:
            pruner (dict): Config for pruner, which is an algorithm component
                for pruning.
        """
        if pruner is None:
            self.pruner = None
            return
        self.pruner = build_pruner(pruner)

        if self.retraining:
            if isinstance(self.channel_cfg, dict):
                self.pruner.deploy_subnet(self.architecture, self.channel_cfg)
                self.deployed = True
            else:
                raise NotImplementedError
        else:
            self.pruner.prepare_from_supernet(self.architecture)

    def _init_distiller(self, distiller):
        """Build registered distillers and make preparations.

        Args:
            distiller (dict): Config for pruner, which is an algorithm
                component for knowledge distillation.
        """
        if distiller is None:
            self.distiller = None
            return
        self.distiller = build_distiller(distiller)
        self.distiller.prepare_from_student(self.architecture)

    @property
    def with_mutator(self):
        """Whether or not this property exists."""
        return hasattr(self, 'mutator') and self.mutator is not None

    @property
    def with_pruner(self):
        """Whether or not this property exists."""
        return hasattr(self, 'pruner') and self.pruner is not None

    @property
    def with_distiller(self):
        """Whether or not this property exists."""
        return hasattr(self, 'distiller') and self.distiller is not None

    def forward(self, img, return_loss=True, **kwargs):
        """Calls either forward_train or forward_test depending on whether
        return_loss=True.

        Note this setting will change the expected inputs. When
        ``return_loss=True``, img and img_meta are single-nested (i.e. Tensor
        and List[dict]), and when ``resturn_loss=False``, img and img_meta
        should be double nested (i.e.  List[Tensor], List[List[dict]]), with
        the outer list indicating test time augmentations.
        """
        return self.architecture(img, return_loss=return_loss, **kwargs)

    def simple_test(self, img, img_metas):
        """Test without augmentation."""
        return self.architecture.simple_test(img, img_metas)

    def show_result(self, img, result, **kwargs):
        """Draw `result` over `img`"""
        return self.architecture.show_result(img, result, **kwargs)

    def _parse_losses(self, losses):
        """Parse the raw outputs (losses) of the network.

        Args:
            losses (dict): Raw output of the network, which usually contain
                losses and other necessary information.
        Returns:
            tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor
                which may be a weighted sum of all losses, log_vars contains
                all the variables to be sent to the logger.
        """
        log_vars = OrderedDict()
        for loss_name, loss_value in losses.items():
            if isinstance(loss_value, torch.Tensor):
                log_vars[loss_name] = loss_value.mean()
            elif isinstance(loss_value, list):
                log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
            elif isinstance(loss_value, dict):
                for name, value in loss_value.items():
                    log_vars[name] = value
            else:
                raise TypeError(
                    f'{loss_name} is not a tensor or list of tensors')

        loss = sum(_value for _key, _value in log_vars.items()
                   if 'loss' in _key)

        log_vars['loss'] = loss
        for loss_name, loss_value in log_vars.items():
            # reduce loss when distributed training
            if dist.is_available() and dist.is_initialized():
                loss_value = loss_value.data.clone()
                dist.all_reduce(loss_value.div_(dist.get_world_size()))
            log_vars[loss_name] = loss_value.item()

        return loss, log_vars

    def train_step(self, data, optimizer):
        """The iteration step during training.

        This method defines an iteration step during training, except for the
        back propagation and optimizer updating, which are done in an optimizer
        hook. Note that in some complicated cases or models, the whole process
        including back propagation and optimizer updating are also defined in
        this method, such as GAN.
        Args:
            data (dict): The output of dataloader.
            optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
                runner is passed to ``train_step()``. This argument is unused
                and reserved.
        Returns:
            dict: It should contain at least 3 keys: ``loss``, ``log_vars``,
                ``num_samples``.
                ``loss`` is a tensor for back propagation, which can be a
                weighted sum of multiple losses.
                ``log_vars`` contains all the variables to be sent to the
                logger.
                ``num_samples`` indicates the batch size (when the model is
                DDP, it means the batch size on each GPU), which is used for
                averaging the logs.
        """
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)

        outputs = dict(
            loss=loss, log_vars=log_vars, num_samples=len(data['img'].data))

        return outputs

    def val_step(self, data, optimizer=None):
        """The iteration step during validation.

        This method shares the same signature as :func:`train_step`, but used
        during val epochs. Note that the evaluation after training epochs is
        not implemented with this method, but an evaluation hook.
        """
        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)

        outputs = dict(
            loss=loss, log_vars=log_vars, num_samples=len(data['img'].data))

        return outputs

mmrazor/models/algorithms/base.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
# Copyright (c) Microsoft Corporation.
import copy

import torch
from torch import nn

from mmrazor.models.builder import ALGORITHMS
from .base import BaseAlgorithm


@ALGORITHMS.register_module()
class Darts(BaseAlgorithm):

    def __init__(self, unroll, **kwargs):

        super(Darts, self).__init__(**kwargs)
        self.unroll = unroll

    def train_step(self, data, optimizer):
        """The iteration step during training.

        This method defines an iteration step during training, except for the
        back propagation and optimizer updating, which are done in an optimizer
        hook. Note that in some complicated cases or models, the whole process
        including back propagation and optimizer updating are also defined in
        this method, such as GAN.

        Args:
            data (dict): The output of dataloader.
            optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
                runner is passed to ``train_step()``. This argument is unused
                and reserved.

        Returns:
            dict: It should contain at least 3 keys: ``loss``, ``log_vars``,
                ``num_samples``.
                ``loss`` is a tensor for back propagation, which can be a
                weighted sum of multiple losses.
                ``log_vars`` contains all the variables to be sent to the
                logger.
                ``num_samples`` indicates the batch size (when the model is
                DDP, it means the batch size on each GPU), which is used for
                averaging the logs.
        """

        if isinstance(data, (tuple, list)) and isinstance(optimizer, dict):
            assert len(data) == len(optimizer)

            train_arch_data, train_supernet_data = data

            optimizer['mutator'].zero_grad()
            if self.unroll:
                self._unrolled_backward(train_arch_data, train_supernet_data,
                                        optimizer)
            else:

                arch_losses = self(**train_arch_data)
                arch_loss, _ = self._parse_losses(arch_losses)
                arch_loss.backward()
            optimizer['mutator'].step()

            model_losses = self(**train_supernet_data)
            model_loss, log_vars = self._parse_losses(model_losses)

            optimizer['architecture'].zero_grad()
            model_loss.backward()
            nn.utils.clip_grad_norm_(
                self.architecture.parameters(), max_norm=5, norm_type=2)
            optimizer['architecture'].step()

            outputs = dict(
                loss=model_loss,
                log_vars=log_vars,
                num_samples=len(train_supernet_data['img'].data))

        else:

            outputs = super(Darts, self).train_step(data, optimizer)
        return outputs

    def _unrolled_backward(self, train_arch_data, train_supernet_data,
                           optimizer):
        """Compute unrolled loss and backward its gradients."""
        backup_params = copy.deepcopy(tuple(self.architecture.parameters()))

        # do virtual step on training data
        lr = optimizer['architecture'].param_groups[0]['lr']
        momentum = optimizer['architecture'].param_groups[0]['momentum']
        weight_decay = optimizer['architecture'].param_groups[0][
            'weight_decay']
        self._compute_virtual_model(train_supernet_data, lr, momentum,
                                    weight_decay, optimizer)

        # calculate unrolled loss on validation data
        # keep gradients for model here for compute hessian
        losses = self(**train_arch_data)
        loss, _ = self._parse_losses(losses)
        w_model, w_arch = tuple(self.architecture.parameters()), tuple(
            self.mutator.parameters())
        w_grads = torch.autograd.grad(loss, w_model + w_arch)
        d_model, d_arch = w_grads[:len(w_model)], w_grads[len(w_model):]

        # compute hessian and final gradients
        hessian = self._compute_hessian(backup_params, d_model,
                                        train_supernet_data)
        with torch.no_grad():
            for param, d, h in zip(w_arch, d_arch, hessian):
                # gradient = dalpha - lr * hessian
                param.grad = d - lr * h

        # restore weights
        self._restore_weights(backup_params)

    def _compute_virtual_model(self, data, lr, momentum, weight_decay,
                               optimizer):
        """Compute unrolled weights w`"""
        # don't need zero_grad, using autograd to calculate gradients
        losses = self(**data)
        loss, _ = self._parse_losses(losses)
        gradients = torch.autograd.grad(loss, self.architecture.parameters())
        with torch.no_grad():
            for w, g in zip(self.architecture.parameters(), gradients):
                m = optimizer['architecture'].state[w].get(
                    'momentum_buffer', 0.)
                w = w - lr * (momentum * m + g + weight_decay * w)

    def _restore_weights(self, backup_params):
        with torch.no_grad():
            for param, backup in zip(self.architecture.parameters(),
                                     backup_params):
                param.copy_(backup)

    def _compute_hessian(self, backup_params, dw, data):
        """
            dw = dw` { L_val(w`, alpha) }
            w+ = w + eps * dw
            w- = w - eps * dw
            hessian = (dalpha { L_trn(w+, alpha) }  \
                - dalpha { L_trn(w-, alpha) }) / (2*eps)
            eps = 0.01 / ||dw||
        """
        self._restore_weights(backup_params)
        norm = torch.cat([w.view(-1) for w in dw]).norm()
        eps = 0.01 / norm
        if norm < 1E-8:
            print(
                'In computing hessian, norm is smaller than 1E-8, \
                cause eps to be %.6f.', norm.item())

        dalphas = []
        for e in [eps, -2. * eps]:
            # w+ = w + eps*dw`, w- = w - eps*dw`
            with torch.no_grad():
                for p, d in zip(self.architecture.parameters(), dw):
                    p += e * d

            losses = self(**data)
            loss, _ = self._parse_losses(losses)
            dalphas.append(
                torch.autograd.grad(loss, tuple(self.mutator.parameters())))
        # dalpha { L_trn(w+) }, # dalpha { L_trn(w-) }
        dalpha_pos, dalpha_neg = dalphas
        hessian = [(p - n) / (2. * eps)
                   for p, n in zip(dalpha_pos, dalpha_neg)]
        return hessian

mmrazor/models/algorithms/darts.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy

from mmcv.cnn import get_model_complexity_info
from mmcv.cnn.utils import revert_sync_batchnorm

from mmrazor.models.builder import ALGORITHMS
from .spos import SPOS


@ALGORITHMS.register_module()
class DetNAS(SPOS):
    """Implementation of `DetNAS <https://arxiv.org/abs/1903.10979>`_"""

    def __init__(self, **kwargs):
        super(DetNAS, self).__init__(**kwargs)

    def _init_flops(self):
        """Get flops of all modules in supernet in order to easily get each
        subnet's flops."""
        flops_model = copy.deepcopy(self.architecture)
        flops_model = revert_sync_batchnorm(flops_model)
        flops_model.eval()
        flops, params = get_model_complexity_info(flops_model.model.backbone,
                                                  self.input_shape)
        flops_lookup = dict()
        for name, module in flops_model.named_modules():
            flops = getattr(module, '__flops__', 0)
            flops_lookup[name] = flops
        del (flops_model)

        for name, module in self.architecture.named_modules():
            module.__flops__ = flops_lookup[name]

mmrazor/models/algorithms/detnas.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmrazor.models.builder import ALGORITHMS
from mmrazor.models.utils import add_prefix
from .base import BaseAlgorithm


@ALGORITHMS.register_module()
class GeneralDistill(BaseAlgorithm):
    """General Distillation Algorithm.

    Args:
        with_student_loss (bool): Whether to use student loss.
            Defaults to True.
        with_teacher_loss (bool): Whether to use teacher loss.
            Defaults to False.
    """

    def __init__(self,
                 with_student_loss=True,
                 with_teacher_loss=False,
                 **kwargs):

        super(GeneralDistill, self).__init__(**kwargs)
        self.with_student_loss = with_student_loss
        self.with_teacher_loss = with_teacher_loss

    def train_step(self, data, optimizer):
        """"""
        losses = dict()
        if self.with_teacher_loss:
            teacher_losses = self.distiller.exec_teacher_forward(data)
            teacher_losses = add_prefix(teacher_losses, 'teacher')
            losses.update(teacher_losses)
        else:
            # Just to be able to trigger the forward hooks that
            # have been registered
            _ = self.distiller.exec_teacher_forward(data)

        if self.with_student_loss:
            student_losses = self.distiller.exec_student_forward(
                self.architecture, data)
            student_losses = add_prefix(student_losses, 'student')
            losses.update(student_losses)
        else:
            # Just to be able to trigger the forward hooks that
            # have been registered
            _ = self.distiller.exec_student_forward(self.architecture, data)

        distill_losses = self.distiller.compute_distill_loss(data)
        distill_losses = add_prefix(distill_losses, 'distiller')
        losses.update(distill_losses)

        loss, log_vars = self._parse_losses(losses)
        outputs = dict(
            loss=loss, log_vars=log_vars, num_samples=len(data['img'].data))
        return outputs

mmrazor/models/algorithms/general_distill.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from functools import partial

from mmcv.cnn import get_model_complexity_info
from torch.nn.modules.batchnorm import _BatchNorm

from mmrazor.models.builder import ALGORITHMS
from .base import BaseAlgorithm


@ALGORITHMS.register_module()
class SPOS(BaseAlgorithm):
    """Implementation of `SPOS <https://arxiv.org/abs/1904.00420>`_"""

    def __init__(self,
                 input_shape=(3, 224, 224),
                 bn_training_mode=False,
                 **kwargs):

        super(SPOS, self).__init__(**kwargs)
        self.input_shape = input_shape
        self.bn_training_mode = bn_training_mode
        if not self.retraining:
            self._init_flops()
        self.apply(partial(self.mutator.reset_in_subnet, in_subnet=True))

    def _init_flops(self):
        """Get flops of all modules in supernet in order to easily get each
        subnet's flops."""
        flops_model = copy.deepcopy(self.architecture)
        flops_model.eval()
        if hasattr(flops_model, 'forward_dummy'):
            flops_model.forward = flops_model.forward_dummy
        else:
            raise NotImplementedError(
                'FLOPs counter is currently not currently supported with {}'.
                format(flops_model.__class__.__name__))

        flops, params = get_model_complexity_info(flops_model,
                                                  self.input_shape)
        flops_lookup = dict()
        for name, module in flops_model.named_modules():
            flops = getattr(module, '__flops__', 0)
            flops_lookup[name] = flops
        del (flops_model)

        for name, module in self.architecture.named_modules():
            module.__flops__ = flops_lookup[name]

    def get_subnet_flops(self):
        """Get subnet's flops based on the complexity information of
        supernet."""
        flops = 0
        for name, module in self.architecture.named_modules():
            if module.__in_subnet__:
                flops += getattr(module, '__flops__', 0)
        return flops

    def train_step(self, data, optimizer):
        """The iteration step during training.

        In retraining stage, to train subnet like common model. In pre-training
        stage, First to sample a subnet from supernet, then to train the
        subnet.
        """
        if self.retraining:
            outputs = super(SPOS, self).train_step(data, optimizer)
        else:
            subnet_dict = self.mutator.sample_subnet()
            self.mutator.set_subnet(subnet_dict)
            outputs = super(SPOS, self).train_step(data, optimizer)
        return outputs

    def train(self, mode=True):
        """Overwrite the train method in `nn.Module` to set `nn.BatchNorm` to
        training mode when model is set to eval mode when
        `self.bn_training_mode` is `True`.

        Args:
            mode (bool): whether to set training mode (`True`) or evaluation
                mode (`False`). Default: `True`.
        """
        super(SPOS, self).train(mode)
        if not mode and self.bn_training_mode:
            for module in self.modules():
                if isinstance(module, _BatchNorm):
                    module.training = True

mmrazor/models/algorithms/spos.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .components import *  # noqa: F401,F403
from .mmcls import MMClsArchitecture
from .mmdet import MMDetArchitecture
from .mmseg import MMSegArchitecture
from .utils import *  # noqa: F401,F403

__all__ = ['MMClsArchitecture', 'MMDetArchitecture', 'MMSegArchitecture']

mmrazor/models/architectures/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.cnn import MODELS
from mmcv.runner import BaseModule


class BaseArchitecture(BaseModule):
    """Base class for architecture.

    Args:
        model (:obj:`torch.nn.Module`): Model to be slimmed, such as
            ``DETECTOR`` in MMDetection.
    """

    def __init__(self, model, **kwargs):
        super(BaseArchitecture, self).__init__(**kwargs)
        self.model = MODELS.build(model)

    def forward_dummy(self, img):
        """Used for calculating network flops."""
        assert hasattr(self.model, 'forward_dummy')
        return self.model.forward_dummy(img)

    def forward(self, img, return_loss=True, **kwargs):
        """Calls either forward_train or forward_test depending on whether
        return_loss=True.

        Note this setting will change the expected inputs. When
        ``return_loss=True``, img and img_meta are single-nested (i.e. Tensor
        and List[dict]), and when ``resturn_loss=False``, img and img_meta
        should be double nested (i.e.  List[Tensor], List[List[dict]]), with
        the outer list indicating test time augmentations.
        """
        return self.model(img, return_loss=return_loss, **kwargs)

    def simple_test(self, img, img_metas):
        """Test without augmentation."""
        return self.model.simple_test(img, img_metas)

    def show_result(self, img, result, **kwargs):
        """Draw `result` over `img`"""
        return self.model.show_result(img, result, **kwargs)

mmrazor/models/architectures/base.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .backbones import *  # noqa: F401,F403
from .heads import *  # noqa: F401,F403
from .necks import *  # noqa: F401,F403

mmrazor/models/architectures/components/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .darts_backbone import DartsBackbone
from .searchable_mobilenet import SearchableMobileNet
from .searchable_shufflenet_v2 import SearchableShuffleNetV2

__all__ = ['DartsBackbone', 'SearchableShuffleNetV2', 'SearchableMobileNet']

mmrazor/models/architectures/components/backbones/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy

import torch
import torch.nn as nn
from mmcls.models.builder import BACKBONES
from mmcv.cnn import build_activation_layer, build_norm_layer

from ...utils import Placeholder


class FactorizedReduce(nn.Module):
    """Reduce feature map size by factorized pointwise (stride=2)."""

    def __init__(self,
                 in_channels,
                 out_channels,
                 act_cfg=dict(type='ReLU'),
                 norm_cfg=dict(type='BN')):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.act_cfg = act_cfg
        self.norm_cfg = norm_cfg
        self.relu = build_activation_layer(self.act_cfg)
        self.conv1 = nn.Conv2d(
            self.in_channels,
            self.out_channels // 2,
            1,
            stride=2,
            padding=0,
            bias=False)
        self.conv2 = nn.Conv2d(
            self.in_channels,
            self.out_channels // 2,
            1,
            stride=2,
            padding=0,
            bias=False)
        self.bn = build_norm_layer(self.norm_cfg, self.out_channels)[1]

    def forward(self, x):
        x = self.relu(x)
        out = torch.cat([self.conv1(x), self.conv2(x[:, :, 1:, 1:])], dim=1)
        out = self.bn(out)
        return out


class StandardConv(nn.Module):
    """
    Standard conv: ReLU - Conv - BN
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 act_cfg=dict(type='ReLU'),
                 norm_cfg=dict(type='BN')):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.act_cfg = act_cfg
        self.norm_cfg = norm_cfg
        self.net = nn.Sequential(
            build_activation_layer(self.act_cfg),
            nn.Conv2d(
                self.in_channels,
                self.out_channels,
                self.kernel_size,
                self.stride,
                self.padding,
                bias=False),
            build_norm_layer(self.norm_cfg, self.out_channels)[1])

    def forward(self, x):
        return self.net(x)


class Node(nn.Module):

    def __init__(self, node_id, num_prev_nodes, channels,
                 num_downsample_nodes):
        super().__init__()
        edges = nn.ModuleDict()
        for i in range(num_prev_nodes):
            if i < num_downsample_nodes:
                stride = 2
            else:
                stride = 1

            edge_id = '{}_p{}'.format(node_id, i)
            edges.add_module(
                edge_id,
                nn.Sequential(
                    Placeholder(
                        group='node',
                        space_id=edge_id,
                        choice_args=dict(
                            stride=stride,
                            in_channels=channels,
                            out_channels=channels)), ))

        self.edges = Placeholder(
            group='node_edge', space_id=node_id, choices=edges)

    def forward(self, prev_nodes):
        return self.edges(prev_nodes)


class Cell(nn.Module):

    def __init__(self,
                 num_nodes,
                 channels,
                 prev_channels,
                 prev_prev_channels,
                 reduction,
                 prev_reduction,
                 act_cfg=dict(type='ReLU'),
                 norm_cfg=dict(type='BN')):
        super().__init__()
        self.act_cfg = act_cfg
        self.norm_cfg = norm_cfg
        self.reduction = reduction
        self.num_nodes = num_nodes

        # If previous cell is reduction cell, current input size does not match
        # with output size of cell[k-2]. So the output[k-2] should be reduced
        # by preprocessing.
        if prev_reduction:
            self.preproc0 = FactorizedReduce(prev_prev_channels, channels,
                                             self.act_cfg, self.norm_cfg)
        else:
            self.preproc0 = StandardConv(prev_prev_channels, channels, 1, 1, 0,
                                         self.act_cfg, self.norm_cfg)
        self.preproc1 = StandardConv(prev_channels, channels, 1, 1, 0,
                                     self.act_cfg, self.norm_cfg)

        # generate dag
        self.nodes = nn.ModuleList()
        for depth in range(2, self.num_nodes + 2):
            if reduction:
                node_id = f'reduce_n{depth}'
                num_downsample_nodes = 2
            else:
                node_id = f'normal_n{depth}'
                num_downsample_nodes = 0
            self.nodes.append(
                Node(node_id, depth, channels, num_downsample_nodes))

    def forward(self, s0, s1):
        # s0, s1 are the outputs of previous previous cell and previous cell,
        # respectively.
        tensors = [self.preproc0(s0), self.preproc1(s1)]
        for node in self.nodes:
            cur_tensor = node(tensors)
            tensors.append(cur_tensor)

        output = torch.cat(tensors[2:], dim=1)
        return output


class AuxiliaryModule(nn.Module):
    """Auxiliary head in 2/3 place of network to let the gradient flow well."""

    def __init__(self,
                 in_channels,
                 base_channels,
                 out_channels,
                 norm_cfg=dict(type='BN')):

        super().__init__()
        self.norm_cfg = norm_cfg
        self.net = nn.Sequential(
            nn.ReLU(),
            nn.AvgPool2d(5, stride=2, padding=0,
                         count_include_pad=False),  # 2x2 out
            nn.Conv2d(in_channels, base_channels, kernel_size=1, bias=False),
            build_norm_layer(self.norm_cfg, base_channels)[1],
            nn.ReLU(inplace=True),
            nn.Conv2d(base_channels, out_channels, kernel_size=2,
                      bias=False),  # 1x1 out
            build_norm_layer(self.norm_cfg, out_channels)[1],
            nn.ReLU(inplace=True))

    def forward(self, x):
        return self.net(x)


@BACKBONES.register_module()
class DartsBackbone(nn.Module):

    def __init__(self,
                 in_channels,
                 base_channels,
                 num_layers=8,
                 num_nodes=4,
                 stem_multiplier=3,
                 out_indices=(7, ),
                 auxliary=False,
                 aux_channels=None,
                 aux_out_channels=None,
                 act_cfg=dict(type='ReLU'),
                 norm_cfg=dict(type='BN')):
        super().__init__()

        self.in_channels = in_channels
        self.base_channels = base_channels
        self.num_layers = num_layers
        self.num_nodes = num_nodes
        self.stem_multiplier = stem_multiplier
        self.out_indices = out_indices
        assert self.out_indices[-1] == self.num_layers - 1
        if auxliary:
            assert aux_channels is not None
            assert aux_out_channels is not None
            self.aux_channels = aux_channels
            self.aux_out_channels = aux_out_channels
            self.auxliary_indice = 2 * self.num_layers // 3

        else:
            self.auxliary_indice = -1
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.out_channels = self.stem_multiplier * self.base_channels
        stem_norm_cfg = copy.deepcopy(self.norm_cfg)
        stem_norm_cfg.update(dict(affine=True))
        self.stem = nn.Sequential(
            nn.Conv2d(
                self.in_channels, self.out_channels, 3, 1, 1, bias=False),
            build_norm_layer(self.norm_cfg, self.out_channels)[1])

        # for the first cell, stem is used for both s0 and s1
        # [!] prev_prev_channels and prev_channels is output channel size,
        # but c_cur is input channel size.
        prev_prev_channels = self.out_channels
        prev_channels = self.out_channels
        self.out_channels = self.base_channels

        self.cells = nn.ModuleList()
        prev_reduction, reduction = False, False
        for i in range(self.num_layers):
            prev_reduction, reduction = reduction, False
            # Reduce featuremap size and double channels in 1/3
            # and 2/3 layer.
            if i == self.num_layers // 3 or i == 2 * self.num_layers // 3:
                self.out_channels *= 2
                reduction = True

            cell = Cell(self.num_nodes, self.out_channels, prev_channels,
                        prev_prev_channels, reduction, prev_reduction,
                        self.act_cfg, self.norm_cfg)
            self.cells.append(cell)

            prev_prev_channels = prev_channels
            prev_channels = self.out_channels * self.num_nodes

            if i == self.auxliary_indice:
                self.auxliary_module = AuxiliaryModule(prev_channels,
                                                       self.aux_channels,
                                                       self.aux_out_channels,
                                                       self.norm_cfg)

    def forward(self, x):
        outs = []
        s0 = s1 = self.stem(x)
        for i, cell in enumerate(self.cells):
            s0, s1 = s1, cell(s0, s1)
            if i in self.out_indices:
                outs.append(s1)
            if i == self.auxliary_indice and self.training:
                aux_feature = self.auxliary_module(s1)

                outs.insert(0, aux_feature)

        return tuple(outs)

mmrazor/models/architectures/components/backbones/darts_backbone.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmcls.models.backbones.base_backbone import BaseBackbone
from mmcls.models.builder import BACKBONES
from mmcls.models.utils import make_divisible
from mmcv.cnn import ConvModule
from torch.nn.modules.batchnorm import _BatchNorm

from ...utils import Placeholder


@BACKBONES.register_module()
class SearchableMobileNet(BaseBackbone):
    """Searchable MobileNet backbone.

    Args:
        first_channels (int): Channel width of first ConvModule. Default: 32.
        last_channels (int): Channel width of last ConvModule. Default: 1200.
        widen_factor (float): Width multiplier, multiply number of
            channels in each layer by this amount. Default: 1.0.
        out_indices (None or Sequence[int]): Output from which stages.
            Default: (7, ).
        frozen_stages (int): Stages to be frozen (all param fixed).
            Default: -1, which means not freezing any parameters.
        conv_cfg (dict, optional): Config dict for convolution layer.
            Default: None, which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='BN').
        act_cfg (dict): Config dict for activation layer.
            Default: dict(type='ReLU6').
        norm_eval (bool): Whether to set norm layers to eval mode, namely,
            freeze running stats (mean and var). Note: Effect on Batch Norm
            and its variants only. Default: False.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.
        arch_setting_type (str): Specify architecture setting.
            Default: 'original'.
        init_cfg (dict | list[dict]): initialization configuration dict to
            define initializer. OpenMMLab has implemented 6 initializers
            including ``Constant``, ``Xavier``, ``Normal``, ``Uniform``,
            ``Kaiming``, and ``Pretrained``.
    """

    # Parameters to build layers. 3 parameters are needed to construct a
    # layer, from left to right: channel, num_blocks, stride.
    arch_settings_dict = {
        'original': [[16, 1, 1], [24, 2, 2], [32, 3, 2], [64, 4, 2],
                     [96, 3, 1], [160, 3, 2], [320, 1, 1]],
        'proxyless_gpu': [[24, 1, 1], [32, 4, 2], [56, 4, 2], [112, 4, 2],
                          [128, 4, 1], [256, 4, 2], [432, 1, 1]],
    }

    def __init__(self,
                 first_channels=32,
                 last_channels=1280,
                 widen_factor=1.,
                 out_indices=(7, ),
                 frozen_stages=-1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU6'),
                 norm_eval=False,
                 with_cp=False,
                 arch_setting_type='original',
                 init_cfg=[
                     dict(type='Kaiming', layer=['Conv2d']),
                     dict(
                         type='Constant',
                         val=1,
                         layer=['_BatchNorm', 'GroupNorm'])
                 ]):
        super(SearchableMobileNet, self).__init__(init_cfg)

        arch_settings = self.arch_settings_dict.get(arch_setting_type)
        if arch_settings is None:
            raise ValueError(f'Expect `arch_setting_type`: '
                             f'{list(self.arch_settings_dict.keys())}, '
                             f'but got: {arch_setting_type}')
        self.arch_settings = arch_settings
        self.widen_factor = widen_factor
        self.out_indices = out_indices
        for index in out_indices:
            if index not in range(0, 8):
                raise ValueError('the item in out_indices must in '
                                 f'range(0, 8). But received {index}')

        if frozen_stages not in range(-1, 8):
            raise ValueError('frozen_stages must be in range(-1, 8). '
                             f'But received {frozen_stages}')
        self.out_indices = out_indices
        self.frozen_stages = frozen_stages
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.norm_eval = norm_eval
        self.with_cp = with_cp

        self.in_channels = make_divisible(first_channels * widen_factor, 8)

        self.conv1 = ConvModule(
            in_channels=3,
            out_channels=self.in_channels,
            kernel_size=3,
            stride=2,
            padding=1,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.layers = []

        for i, layer_cfg in enumerate(self.arch_settings):
            channel, num_blocks, stride = layer_cfg
            out_channels = make_divisible(channel * widen_factor, 8)
            inverted_res_layer = self.make_layer(
                out_channels=out_channels,
                num_blocks=num_blocks,
                stride=stride,
                stage_idx=i)
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, inverted_res_layer)
            self.layers.append(layer_name)

        if widen_factor > 1.0:
            self.out_channel = int(last_channels * widen_factor)
        else:
            self.out_channel = last_channels

        layer = ConvModule(
            in_channels=self.in_channels,
            out_channels=self.out_channel,
            kernel_size=1,
            stride=1,
            padding=0,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        self.add_module('conv2', layer)
        self.layers.append('conv2')

    def make_layer(self, out_channels, num_blocks, stride, stage_idx):
        """Stack InvertedResidual blocks to build a layer for MobileNetV2.

        Args:
            out_channels (int): out_channels of block.
            num_blocks (int): number of blocks.
            stride (int): stride of the first block. Default: 1
            expand_ratio (int): Expand the number of channels of the
                hidden layer in InvertedResidual by this ratio. Default: 6.
        """
        layers = []
        for i in range(num_blocks):
            if i >= 1:
                stride = 1
            # HACK
            # do not search first block
            if stage_idx == 0:
                group = 'first_blocks'
            else:
                group = 'searchable_blocks'
            layers.append(
                Placeholder(
                    group=group,
                    space_id=f'stage_{stage_idx}_block_{i}',
                    choice_args=dict(
                        in_channels=self.in_channels,
                        out_channels=out_channels,
                        stride=stride,
                    )))
            self.in_channels = out_channels

        return nn.Sequential(*layers)

    def forward(self, x):
        """Forward computation.

        Args:
            x (tensor | tuple[tensor]): x could be a Torch.tensor or a tuple of
                Torch.tensor, containing input data for forward computation.
        """
        x = self.conv1(x)

        outs = []
        for i, layer_name in enumerate(self.layers):
            layer = getattr(self, layer_name)
            x = layer(x)
            if i in self.out_indices:
                outs.append(x)

        return tuple(outs)

    def _freeze_stages(self):
        """Freeze params not to update in the specified stages."""
        if self.frozen_stages >= 0:
            for param in self.conv1.parameters():
                param.requires_grad = False
        for i in range(1, self.frozen_stages + 1):
            layer = getattr(self, f'layer{i}')
            layer.eval()
            for param in layer.parameters():
                param.requires_grad = False

    def train(self, mode=True):
        """Set module status before forward computation.

        Args:
            mode (bool): Whether it is train_mode or test_mode
        """
        super(SearchableMobileNet, self).train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, _BatchNorm):
                    m.eval()

mmrazor/models/architectures/components/backbones/searchable_mobilenet.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmcls.models.backbones.base_backbone import BaseBackbone
from mmcls.models.builder import BACKBONES
from mmcv.cnn import ConvModule, constant_init, normal_init
from torch.nn.modules.batchnorm import _BatchNorm

from ...utils import Placeholder


@BACKBONES.register_module()
class SearchableShuffleNetV2(BaseBackbone):
    """Based on ShuffleNetV2 backbone.
    Args:
        widen_factor (float): Width multiplier - adjusts the number of
            channels in each layer by this amount. Default: 1.0.
        out_indices (Sequence[int]): Output from which stages.
            Default: (0, 1, 2, 3).
        frozen_stages (int): Stages to be frozen (all param fixed).
            Default: -1, which means not freezing any parameters.
        with_last_layer (bool): Whether is last layer.
            Default: True, which means not need to add `Placeholder``.
        conv_cfg (dict, optional): Config dict for convolution layer.
            Default: None, which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='BN').
        act_cfg (dict): Config dict for activation layer.
            Default: dict(type='ReLU').
        norm_eval (bool): Whether to set norm layers to eval mode, namely,
            freeze running stats (mean and var). Note: Effect on Batch Norm
            and its variants only. Default: False.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.
        init_cfg (dict): Init config dict for ``BaseBackbone``.
    """

    def __init__(self,
                 stem_multiplier=1,
                 widen_factor=1.0,
                 out_indices=(4, ),
                 frozen_stages=-1,
                 with_last_layer=True,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 norm_eval=False,
                 with_cp=False,
                 **kwargs):
        super(SearchableShuffleNetV2, self).__init__(**kwargs)
        self.stage_blocks = [4, 4, 8, 4]
        self.widen_factor = widen_factor
        self.out_indices = out_indices
        self.frozen_stages = frozen_stages
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.norm_eval = norm_eval
        self.with_cp = with_cp

        channels = [
            round(64 * self.widen_factor),
            round(160 * self.widen_factor),
            round(320 * self.widen_factor),
            round(640 * self.widen_factor)
        ]
        last_channels = 1024

        self.in_channels = 16 * stem_multiplier
        self.conv1 = ConvModule(
            in_channels=3,
            out_channels=self.in_channels,
            kernel_size=3,
            stride=2,
            padding=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)

        self.layers = nn.ModuleList()
        for i, num_blocks in enumerate(self.stage_blocks):
            layer = self._make_layer(channels[i], num_blocks, i)
            self.layers.append(layer)

        if with_last_layer:
            self.layers.append(
                ConvModule(
                    in_channels=self.in_channels,
                    out_channels=last_channels,
                    kernel_size=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg))

    def _make_layer(self, out_channels, num_blocks, stage_idx):
        """Stack blocks to make a layer.

        Args:
            out_channels (int): out_channels of the block.
            num_blocks (int): number of blocks.
            stage_idx (int): used in ``space_id``.
        Returns:
            torch.nn.Sequential: The layer made.
        """
        layers = []
        for i in range(num_blocks):
            stride = 2 if i == 0 else 1
            layers.append(
                Placeholder(
                    group='all_blocks',
                    space_id=f'stage_{stage_idx}_block_{i}',
                    choice_args=dict(
                        in_channels=self.in_channels,
                        out_channels=out_channels,
                        stride=stride,
                    )))
            self.in_channels = out_channels

        return nn.Sequential(*layers)

    def _freeze_stages(self):
        """Freeze params not to update in the specified stages."""
        if self.frozen_stages >= 0:
            for param in self.conv1.parameters():
                param.requires_grad = False

        for i in range(self.frozen_stages):
            m = self.layers[i]
            m.eval()
            for param in m.parameters():
                param.requires_grad = False

    def init_weights(self):
        """Init weights of ``SearchableShuffleNetV2``."""
        super(SearchableShuffleNetV2, self).init_weights()

        if (isinstance(self.init_cfg, dict)
                and self.init_cfg['type'] == 'Pretrained'):
            # Suppress default init if use pretrained model.
            return

        for name, m in self.named_modules():
            if isinstance(m, nn.Conv2d):
                if 'conv1' in name:
                    normal_init(m, mean=0, std=0.01)
                else:
                    normal_init(m, mean=0, std=1.0 / m.weight.shape[1])
            elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
                constant_init(m, val=1, bias=0.0001)
                if isinstance(m, _BatchNorm):
                    if m.running_mean is not None:
                        nn.init.constant_(m.running_mean, 0)

    def forward(self, x):
        """Forward computation.

        Args:
            x (tensor | tuple[tensor]): x could be a Torch.tensor or a tuple of
                Torch.tensor, containing input data for forward computation.
        """
        x = self.conv1(x)

        outs = []
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i in self.out_indices:
                outs.append(x)

        return tuple(outs)

    def train(self, mode=True):
        """Set module status before forward computation.

        Args:
            mode (bool): Whether it is train_mode or test_mode
        """
        super(SearchableShuffleNetV2, self).train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, _BatchNorm):
                    m.eval()

mmrazor/models/architectures/components/backbones/searchable_shufflenet_v2.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .darts_head import DartsSubnetClsHead

__all__ = ['DartsSubnetClsHead']

mmrazor/models/architectures/components/heads/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.

from mmcls.models import HEADS, build_loss
from mmcls.models.heads import LinearClsHead
from torch import nn

from mmrazor.models.utils import add_prefix


@HEADS.register_module()
class DartsSubnetClsHead(LinearClsHead):

    def __init__(self, aux_in_channels, aux_loss, **kwargs):
        super(DartsSubnetClsHead, self).__init__(**kwargs)
        self.aux_linear = nn.Linear(aux_in_channels, self.num_classes)
        self.compute_aux_loss = build_loss(aux_loss)

    def forward_train(self, x, gt_label, **kwargs):

        losses = super().forward_train(x, gt_label, **kwargs)
        losses = add_prefix(losses, 'subnet')

        assert isinstance(x, tuple) and len(x) > 1
        x = x[0]

        cls_score = self.aux_linear(x)
        num_samples = len(cls_score)
        aux_loss = self.compute_aux_loss(
            cls_score, gt_label, avg_factor=num_samples, **kwargs)

        losses['aux_head.loss'] = aux_loss

        return losses

mmrazor/models/architectures/components/heads/darts_head.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmcls.models.builder import HEADS
from mmcls.models.heads import LinearClsHead
from torch import nn


@HEADS.register_module()
class LinearNoBiasClsHead(LinearClsHead):
    """Set bias to False in ``LinearClsHead``."""

    def __init__(self, **kwargs):
        super(LinearNoBiasClsHead, self).__init__(**kwargs)
        self.fc = nn.Linear(self.in_channels, self.num_classes, bias=False)

mmrazor/models/architectures/components/heads/no_bias_fc_head.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from ..builder import ARCHITECTURES
from .base import BaseArchitecture


@ARCHITECTURES.register_module()
class MMClsArchitecture(BaseArchitecture):
    """Architecture based on MMCls."""

    def __init__(self, **kwargs):
        super(MMClsArchitecture, self).__init__(**kwargs)

    def forward_dummy(self, img):
        """Used for calculating network flops."""
        output = img
        for name, child in self.model.named_children():
            if name == 'head':
                output = child.fc(output[0])
            else:
                output = child(output)
        return output

    def cal_pseudo_loss(self, pseudo_img):
        """Used for executing ``forward`` with pseudo_img."""
        return sum(pseudo_img)

mmrazor/models/architectures/mmcls.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.

from ..builder import ARCHITECTURES
from .base import BaseArchitecture


@ARCHITECTURES.register_module()
class MMDetArchitecture(BaseArchitecture):
    """Architecture based on MMDet."""

    def __init__(self, **kwargs):
        super(MMDetArchitecture, self).__init__(**kwargs)

    def cal_pseudo_loss(self, pseudo_img):
        """Used for executing ``forward`` with pseudo_img."""
        out = 0.
        for levels in pseudo_img:
            out += sum([level.sum() for level in levels])

        return out

mmrazor/models/architectures/mmdet.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .placeholder import Placeholder

__all__ = ['Placeholder']

mmrazor/models/architectures/utils/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from torch import nn


class Placeholder(nn.Module):
    """Used for build searchable network.

    Args:
        group (str): Placeholder_group, as group name in searchable network.
        space_id (str): It is one and only index for each ``Placeholder``.
        choices (dict): Consist of the registered ``OPS``, used to combine
            ``MUTABLES``, the ``Placeholder`` will be replace with the
            ``MUTABLES``.
        choice_args (dict): The configuration of ``OPS`` used in choices.
    """

    def __init__(self, group, space_id, choices=None, choice_args=None):
        super(Placeholder, self).__init__()
        self.placeholder_group = group
        self.placeholder_kwargs = dict(space_id=space_id)
        if choices is not None:
            self.placeholder_kwargs.update(dict(choices=choices))
        if choice_args is not None:
            self.placeholder_kwargs.update(dict(choice_args=choice_args))

mmrazor/models/architectures/utils/placeholder.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.cnn import MODELS as MMCV_MODELS
from mmcv.utils import Registry

MODELS = Registry('models', parent=MMCV_MODELS)

ALGORITHMS = MODELS
MUTABLES = MODELS
DISTILLERS = MODELS
LOSSES = MODELS
OPS = MODELS
PRUNERS = MODELS
QUANTIZERS = MODELS
ARCHITECTURES = MODELS
MUTATORS = MODELS


def build_algorithm(cfg):
    """Build compressor."""
    return ALGORITHMS.build(cfg)


def build_architecture(cfg):
    """Build architecture."""
    return ARCHITECTURES.build(cfg)


def build_mutator(cfg):
    """Build mutator."""
    return MUTATORS.build(cfg)


def build_distiller(cfg):
    """Build distiller."""
    return DISTILLERS.build(cfg)


def build_pruner(cfg):
    """Build pruner."""
    return PRUNERS.build(cfg)


def build_mutable(cfg):
    """Build mutable."""
    return MUTABLES.build(cfg)


def build_op(cfg):
    """Build op."""
    return OPS.build(cfg)


def build_loss(cfg):
    """Build loss."""
    return LOSSES.build(cfg)

mmrazor/models/builder.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .self_distiller import SelfDistiller
from .single_teacher import SingleTeacherDistiller

__all__ = ['SelfDistiller', 'SingleTeacherDistiller']

mmrazor/models/distillers/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod

from mmcv.runner import BaseModule
from mmcv.utils import import_modules_from_strings


def function_wrapper(ctx, method, method_str):
    """Pass teacher's outputs to student."""

    def wrapper(*args, **kwargs):
        # record inputs
        ctx.method_args[method_str] = args
        ctx.method_kwargs[method_str] = kwargs
        # TODO cover more usecases, not only pass teacher's outputs to
        # student.
        if ctx.is_teacher:
            # execute the raw function
            outputs = method(*args, **kwargs)
            # record outputs
            ctx.method_return[method_str] = outputs
        else:
            # modify student's outputs to be same with teacher
            outputs = ctx.method_return[method_str]

        return outputs

    return wrapper


class FunctionContext():
    """Function context manager for rewrite function.

    Args:
        ctx (ConversionContext): The distiller's overall context manager.
        method (str): The name of the function to rewrite.
    """

    def __init__(self, ctx, method, import_module=None):
        self.ctx = ctx

        self.import_module = import_modules_from_strings(import_module)
        self.method_str = method
        self.method_exec_str = f'self.import_module.{method}'

    def _set_method(self, method):
        """Modify a function."""
        exec(f'{self.method_exec_str} = method')

    def __enter__(self):
        """Rewrite the function."""
        self.method_impl = eval(self.method_exec_str)

        if self.method_impl:
            self._set_method(
                function_wrapper(self.ctx, self.method_impl, self.method_str,
                                 self.align_mode))

    def __exit__(self, exc_type, exc_value, traceback):
        """Restore the function."""
        if self.method_impl:
            self._set_method(self.method_impl)


class ConversionContext():
    """Context manager for record functions' inputs or outputs."""

    def __init__(self, hooks):
        # save functions' inputs
        self.method_args = dict()
        self.method_kwargs = dict()
        # save functions' outputs
        self.method_return = dict()

        # Each function will have a sub context manager, the function will be
        # rewritten when enter the sub context manager.
        self.hooks = []
        self.is_teacher = True
        for hook in hooks:
            self.hooks.append(FunctionContext(self, **hook))

    def __enter__(self):
        """Enter every sub context managers."""
        for hook in self.hooks:
            hook.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Exit every sub context managers."""
        for hook in self.hooks:
            hook.__exit__(exc_type, exc_value, traceback)


class BaseDistiller(BaseModule, metaclass=ABCMeta):
    """Base Distiller.

    In the distillation algorithm, some intermediate results of the teacher
    need to be obtained and passed to the student.

    For nn.Module's outputs, obtained by pytorch forward hook.
    For python function's outputs, obtained by a specific context manager.

    Args:
        align_methods (dict): The details of the functions which outputs need
        to be obtained.
    """

    def __init__(self, align_methods=None, **kwargs):
        super(BaseDistiller, self).__init__(**kwargs)

        if align_methods is None:
            self.context_manager = None
        else:
            # To obtain the python function's outputs, there will build a
            # specific context manager. When enter the context manager, the
            # functions will be rewrite. The context manager could record
            # inputs or outputs of the functions , and pass from teachr to
            # student. When exit the context manager, the rewritten functions
            # will restore.
            self.context_manager = ConversionContext(align_methods)

    @abstractmethod
    def prepare_from_student(self, student):
        """Register forward hooks to students and teachers."""
        pass

    @abstractmethod
    def teacher_forward_output_hook(self, module, inputs, outputs):
        """Save the teacher output."""
        pass

    @abstractmethod
    def student_forward_output_hook(self, module, inputs, outputs):
        """Save the student output."""
        pass

    def reset_ctx_teacher_mode(self, mode=True):
        if self.context_manager is not None:
            self.context_manager.is_teacher = mode

    @abstractmethod
    def exec_teacher_forward(self, data):
        """Execute the teacher's forward function."""
        pass

    @abstractmethod
    def exec_student_forward(self, student, data):
        """Execute the student's forward function."""
        pass

    @abstractmethod
    def compute_distill_loss(self, data):
        """Compute distill loss according teacher's outputs and student's
        outputs."""
        pass

mmrazor/models/distillers/base.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn

from ..builder import DISTILLERS, build_loss
from .base import BaseDistiller


@DISTILLERS.register_module()
class SelfDistiller(BaseDistiller):
    """Transfer knowledge inside a single model.

    Args:
        components (dict): The details of the distillation. It usually includes
            the module names of the teacher and the student, and the losses
            used in the distillation.
    """

    def __init__(self, components, **kwargs):
        super().__init__(**kwargs)
        self.components = components
        self.losses = nn.ModuleDict()

        self.student_outputs = dict()
        self.teacher_outputs = dict()

        for component in self.components:
            student_module_name = component['student_module']
            teacher_module_name = component['teacher_module']
            self.student_outputs[student_module_name] = list()
            self.teacher_outputs[teacher_module_name] = list()

            for loss in component.losses:
                loss_cfg = loss.copy()
                loss_name = loss_cfg.pop('name')
                self.losses[loss_name] = build_loss(loss_cfg)

    def prepare_from_student(self, student):
        """Registers a global forward hook for each teacher module and student
        module to be used in the distillation.

        Args:
            student (:obj:`torch.nn.Module`): The student model to be used
                in the distillation.
        """
        self.module2name = {}
        for name, module in student.model.named_modules():
            self.module2name[module] = name
        self.name_modules = dict(student.model.named_modules())

        for component in self.components:
            student_module_name = component['student_module']
            teacher_module_name = component['teacher_module']

            student_module = self.name_modules[student_module_name]
            teacher_module = self.name_modules[teacher_module_name]

            student_module.register_forward_hook(
                self.student_forward_output_hook)
            teacher_module.register_forward_hook(
                self.teacher_forward_output_hook)

    def teacher_forward_output_hook(self, module, inputs, outputs):
        """Save the output.

        Args:
            module (:obj:`torch.nn.Module`): the module of register hook
            inputs (tuple): input of module
            outputs (tuple): out of module
        """
        if self.training and getattr(self, 'is_teacher', None):
            self.teacher_outputs[self.module2name[module]].append(outputs)

    def student_forward_output_hook(self, module, inputs, outputs):
        """Save the output.

        Args:
            module (:obj:`torch.nn.Module`): the module of register hook
            inputs (tuple): input of module
            outputs (tuple): out of module
        """
        if self.training and not getattr(self, 'is_teacher', None):
            self.student_outputs[self.module2name[module]].append(outputs)

    def reset_outputs(self, outputs):
        """Reset the teacher's outputs or student's outputs."""
        for key in outputs.keys():
            outputs[key] = list()

    def exec_teacher_forward(self, teacher, data):
        """Forward computation of the teacher.

        Args:
            teacher (:obj:`torch.nn.Module`): The teacher model to be used
                in the distillation.
            data (dict): The output of dataloader.
        """
        self.reset_outputs(self.teacher_outputs)
        self.is_teacher = True
        output = teacher(**data)
        self.is_teacher = False

        return output

    def exec_student_forward(self, student, data):
        """Forward computation of the student.

        Args:
            student (:obj:`torch.nn.Module`): The student model to be used
                in the distillation.
            data (dict): The output of dataloader.
        """
        assert not self.is_teacher
        self.reset_outputs(self.student_outputs)
        output = student(**data)

        return output

    def compute_distill_loss(self, data):
        """Compute the distillation loss."""

        losses = dict()

        for i, component in enumerate(self.components):
            student_module_name = component['student_module']
            student_outputs = self.student_outputs[student_module_name]

            teacher_module_name = component['teacher_module']
            teacher_outputs = self.teacher_outputs[teacher_module_name]

            for out_idx, (s_out, t_out) in enumerate(
                    zip(student_outputs, teacher_outputs)):

                for loss in component.losses:
                    loss_module = self.losses[loss.name]
                    loss_name = f'{loss.name}.{out_idx}'

                    loss_module.current_data = data
                    losses[loss_name] = loss_module(s_out, t_out)
                    loss_module.current_data = None

        return losses

mmrazor/models/distillers/self_distiller.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from torch.nn.modules.batchnorm import _BatchNorm

from ..builder import DISTILLERS, MODELS, build_loss
from .base import BaseDistiller


@DISTILLERS.register_module()
class SingleTeacherDistiller(BaseDistiller):
    """Distiller with single teacher.

    Args:
        teacher (dict): The config dict for teacher.
        teacher_trainable (bool): Whether the teacher is trainable.
            Default: False.
        teacher_norm_eval (bool): Whether to set teacher's norm layers to eval
            mode, namely, freeze running stats (mean and var). Note: Effect on
            Batch Norm and its variants only. Default: True.
        components (dict): The details of the distillation. It usually includes
            the module names of the teacher and the student, and the losses
            used in the distillation.
    """

    def __init__(self,
                 teacher,
                 teacher_trainable=False,
                 teacher_norm_eval=True,
                 components=tuple(),
                 **kwargs):
        super().__init__(**kwargs)
        self.teacher_trainable = teacher_trainable
        self.teacher_norm_eval = teacher_norm_eval
        self.teacher = self.build_teacher(teacher)

        self.components = components
        self.losses = nn.ModuleDict()
        self.align_modules = nn.ModuleDict()

        # Record the featuremaps that need to calculate the distillation loss.
        self.student_outputs = dict()
        self.teacher_outputs = dict()

        for i, component in enumerate(self.components):
            student_module_name = component['student_module']
            teacher_module_name = component['teacher_module']
            # The type of every student_output is a list by default, because
            # some modules will execute multiple forward calculations, such as
            # the shareable head in Retinanet
            self.student_outputs[student_module_name] = list()
            self.teacher_outputs[teacher_module_name] = list()

            # If the number of featuremap channels of student and teacher are
            # inconsistent, they need to be aligned by a 1x1 convolution
            align_module_cfg = getattr(component, 'align_module', None)
            if align_module_cfg is not None:
                align_module_name = f'component_{i}'
                align_module = self.build_align_module(align_module_cfg)
                self.align_modules[align_module_name] = align_module

            # Multiple losses can be calculated at the same location
            for loss in component.losses:
                loss_cfg = loss.copy()
                loss_name = loss_cfg.pop('name')
                self.losses[loss_name] = build_loss(loss_cfg)

    def build_teacher(self, cfg):
        """Build a model from the `cfg`."""

        teacher = MODELS.build(cfg)

        return teacher

    def build_align_module(self, cfg):
        """Build ``align_module`` from the `cfg`.

        ``align_module`` is needed when the number of channels output by the
        teacher module is not equal to that of the student module, or for some
        other reasons.

        Args:
            cfg (dict): The config dict for ``align_module``.
        """

        in_channels = cfg.student_channels
        out_channels = cfg.teacher_channels
        if cfg.type == 'conv2d':
            align_module = nn.Conv2d(in_channels, out_channels, 1)
        elif cfg.type == 'linear':
            align_module = nn.Linear(in_channels, out_channels)
        return align_module

    def prepare_from_student(self, student):
        """Registers a global forward hook for each teacher module and student
        module to be used in the distillation.

        Args:
            student (:obj:`torch.nn.Module`): The student model to be used
                in the distillation.
        """

        # Record the mapping relationship between student's modules and module
        # names.
        self.student_module2name = {}
        for name, module in student.model.named_modules():
            self.student_module2name[module] = name
        self.student_name2module = dict(student.model.named_modules())

        # Record the mapping relationship between teacher's modules and module
        # names.
        self.teacher_module2name = {}
        for name, module in self.teacher.named_modules():
            self.teacher_module2name[module] = name
        self.teacher_name2module = dict(self.teacher.named_modules())

        # Register forward hooks for modules that need to participate in loss
        # calculation.
        for component in self.components:
            student_module_name = component['student_module']
            teacher_module_name = component['teacher_module']

            student_module = self.student_name2module[student_module_name]
            teacher_module = self.teacher_name2module[teacher_module_name]

            student_module.register_forward_hook(
                self.student_forward_output_hook)
            teacher_module.register_forward_hook(
                self.teacher_forward_output_hook)

    def teacher_forward_output_hook(self, module, inputs, outputs):
        """Save the module's forward output.

        Args:
            module (:obj:`torch.nn.Module`): The module to register hook.
            inputs (tuple): The input of the module.
            outputs (tuple): The output of the module.
        """
        if self.training:
            self.teacher_outputs[self.teacher_module2name[module]].append(
                outputs)

    def student_forward_output_hook(self, module, inputs, outputs):
        """Save the module's forward output.

        Args:
            module (:obj:`torch.nn.Module`): The module to register hook.
            inputs (tuple): The input of the module.
            outputs (tuple): The output of the module.
        """
        if self.training:
            self.student_outputs[self.student_module2name[module]].append(
                outputs)

    def reset_outputs(self, outputs):
        """Reset the teacher's outputs or student's outputs."""
        for key in outputs.keys():
            outputs[key] = list()

    def exec_teacher_forward(self, data):
        """Execute the teacher's forward function.

        After this function, the teacher's featuremaps will be saved in
        ``teacher_outputs``.
        """

        # Convert the context manager's mode to teacher.
        self.reset_ctx_teacher_mode(True)
        # Clear the saved data of the last forward。
        self.reset_outputs(self.teacher_outputs)

        if self.teacher_trainable:
            output = self.teacher(**data)
        else:
            with torch.no_grad():
                output = self.teacher(**data)

        return output

    def exec_student_forward(self, student, data):
        """Execute the teacher's forward function.

        After this function, the student's featuremaps will be saved in
        ``student_outputs``.
        """
        # Convert the context manager's mode to teacher.
        self.reset_ctx_teacher_mode(False)
        # Clear the saved data of the last forward。
        self.reset_outputs(self.student_outputs)

        output = student(**data)
        return output

    def train(self, mode=True):
        """Set distiller's forward mode."""
        super(SingleTeacherDistiller, self).train(mode)
        if mode and self.teacher_norm_eval:
            for m in self.teacher.modules():
                if isinstance(m, _BatchNorm):
                    m.eval()

    def get_teacher_outputs(self, teacher_module_name):
        """Get the outputs according module name."""
        return self.teacher_outputs[teacher_module_name]

    def compute_distill_loss(self, data=None):
        """Compute the distillation loss."""

        losses = dict()

        for i, component in enumerate(self.components):
            # Get the student's outputs.
            student_module_name = component['student_module']
            student_outputs = self.student_outputs[student_module_name]

            # Align student output's channels with teacher.
            align_module_name = f'component_{i}'
            if align_module_name in self.align_modules:
                align_module = self.align_modules[align_module_name]
                student_outputs = [
                    align_module(s_out) for s_out in student_outputs
                ]

            # Get the teacher's outputs.
            teacher_module_name = component['teacher_module']
            teacher_outputs = self.get_teacher_outputs(teacher_module_name)

            # One module maybe have N outputs, such as the shareable head in
            # RetinaNet.
            for out_idx, (s_out, t_out) in enumerate(
                    zip(student_outputs, teacher_outputs)):

                for loss in component.losses:
                    loss_module = self.losses[loss.name]
                    loss_name = f'{loss.name}.{out_idx}'
                    # TODO ugly implementation.
                    # Pass the gt_label to loss function.
                    # Only used by WSLD.
                    loss_module.current_data = data
                    losses[loss_name] = loss_module(s_out, t_out)
                    loss_module.current_data = None

        return losses

mmrazor/models/distillers/single_teacher.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .cwd import ChannelWiseDivergence
from .kl_divergence import KLDivergence
from .relational_kd import AngleWiseRKD, DistanceWiseRKD
from .weighted_soft_label_distillation import WSLD
from .akd_cwd import AKD_ChannelWiseDivergence
from .akd_maskd import AKD_MasKDLoss
from .akd_mgd import AKD_MGDLoss


__all__ = [
    'ChannelWiseDivergence', 'KLDivergence', 'AngleWiseRKD', 'DistanceWiseRKD',
    'WSLD', 'AKD_ChannelWiseDivergence', 'AKD_MasKDLoss', 'AKD_MGDLoss'
]

mmrazor/models/losses/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
import torch.nn.functional as F

from ..builder import LOSSES


@LOSSES.register_module()
class AKD_ChannelWiseDivergence(nn.Module):
    """PyTorch version of `Channel-wise Distillation for Semantic Segmentation.

    <https://arxiv.org/abs/2011.13256>`_.

    Args:
        tau (float): Temperature coefficient. Defaults to 1.0.
        loss_weight (float): Weight of loss. Defaults to 1.0.
    """

    def __init__(
        self,
        tau=1.0,
        loss_weight=1.0,
        drop_rate=0.1
    ):
        super(AKD_ChannelWiseDivergence, self).__init__()
        self.tau = tau
        self.loss_weight = loss_weight

        self.Uncertain = nn.BatchNorm2d(2048, affine=False)
        self.drop_rate = drop_rate

    def forward(self, y_s, y_t):
        """Forward computation.
        Args:
            y_s (list): The student model prediction with
                shape (N, C, H, W) in list.
            y_t (list): The teacher model prediction with
                shape (N, C, H, W) in list.
        Return:
            torch.Tensor: The calculated loss value of all stages.
        """
        if not isinstance(y_s, (tuple, list)):
            y_s = (y_s, )
            y_t = (y_t, )
        assert len(y_s) == len(y_t)
        losses = []

        for idx, (s, t) in enumerate(zip(y_s, y_t)):
            assert s.shape == t.shape

            N, C, H, W = s.shape

            # 1. multi-tea
            multi_T = F.dropout(t, p=self.drop_rate, training=True)

            # 2. uncertainty weights
            s = self.Uncertain(s)
            t = self.Uncertain(t)

            # normalize in channel diemension
            softmax_pred_T = F.softmax(multi_T.view(-1, W * H) / self.tau, dim=1) # [N*C, H*W]

            logsoftmax = torch.nn.LogSoftmax(dim=1)
            cost = torch.sum(softmax_pred_T *
                            logsoftmax(multi_T.view(-1, W * H) / self.tau) -
                            softmax_pred_T *
                            logsoftmax(s.view(-1, W * H) / self.tau)) * (
                                self.tau**2)
            
            losses.append(cost / (C * N))
        loss = sum(losses)

        return loss * self.loss_weight

mmrazor/models/losses/akd_cwd.py
Gumpest/AvatarKD
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from mmcv.runner.checkpoint import _load_checkpoint

from ..builder import LOSSES


def dice_coeff(inputs):
    # inputs: [B, T, H*W]
    pred = inputs[:, None, :, :]
    target = inputs[:, :, None, :]

    mask = pred.new_ones(pred.size(0), target.size(1), pred.size(2))
    mask[:, torch.arange(mask.size(1)), torch.arange(mask.size(2))] = 0

    a = torch.sum(pred * target, -1)
    b = torch.sum(pred * pred, -1) + 1e-12
    c = torch.sum(target * target, -1) + 1e-12
    d = (2 * a) / (b + c)
    d = (d * mask).sum() / mask.sum()
    return d


class MaskModule(nn.Module):

    def __init__(self, channels, num_tokens=8, weight_mask=False):
        super().__init__()
        self.weight_mask = weight_mask
        self.mask_token = nn.Parameter(torch.randn(num_tokens, channels).normal_(0, 0.01))
        if self.weight_mask:
            self.prob = nn.Sequential(
                nn.Conv2d(channels, channels, kernel_size=3, padding=1),
                nn.AdaptiveAvgPool2d(1),
                nn.Conv2d(channels, num_tokens, kernel_size=1)
            )
            self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels  # fan-out
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
                if m.bias is not None:
                    m.bias.data.zero_()

    def forward_mask(self, x):
        N, C, H, W = x.shape
        mask_token = self.mask_token.expand(N, -1, -1)  # [N, T, C]
        k = x.view(N, -1, H * W)
        attn = mask_token @ k  # [N, T, H * W]
        attn = attn.sigmoid()
        attn = attn.view(N, -1, H, W)
        return attn

    def forward_prob(self, x):
        mask_probs = self.prob(x)  # [N, T, 1, 1]
        mask_probs = mask_probs.softmax(1).unsqueeze(2)  # [N, T, 1, 1, 1]
        return mask_probs

    def forward_train(self, x):
        mask = self.forward_mask(x)
        out = x.unsqueeze(1) * mask.unsqueeze(2)  # [N, T, C, H, W]
        # probs
        if self.weight_mask:
            mask_probs = self.forward_prob(x)
            #print(mask_probs.detach().flatten(1))
            out = out * mask_probs
        out = out.sum(1)
        # loss
        mask_loss = dice_coeff(mask.flatten(2))
        return out, mask_loss

    def forward(self, x):
        return self.forward_train(x)


@LOSSES.register_module()
class AKD_MasKDLoss(nn.Module):

    def __init__(self, channels, num_tokens=6, weight_mask=True, custom_mask=True, custom_mask_warmup=1000, pretrained=None, loss_weight=1., drop_rate=0.1):
        super().__init__()
        self.loss_weight = loss_weight
        self.weight_mask = weight_mask
        self.custom_mask = custom_mask
        self.custom_mask_warmup = custom_mask_warmup
        self.drop_rate = drop_rate

        self.mask_modules = nn.ModuleList([
            MaskModule(channels=c, num_tokens=num_tokens, weight_mask=weight_mask) for c in channels]
        )

        self.init_weights(pretrained)
        self._iter = 0

        self.Uncertain = nn.ModuleList([
            nn.BatchNorm2d(c, affine=False) for c in channels]
        )

    def init_weights(self, pretrained=None):
        if pretrained is None:
            return
        ckpt = _load_checkpoint(pretrained, map_location='cpu')
        state_dict = {}
        for k, v in ckpt['state_dict'].items():
            if 'mask_modules' in k:
                state_dict[k[k.find('mask_modules'):]] = v
        self.load_state_dict(state_dict, strict=True)

    def forward(self, y_s_list, y_t_list):
        if not isinstance(y_s_list, (tuple, list)):
            y_s_list = (y_s_list, )
            y_t_list = (y_t_list, )
        assert len(y_s_list) == len(y_t_list) == len(self.mask_modules)

        losses = []
        for y_s, y_t, Uncertain, mask_module in zip(y_s_list, y_t_list, self.Uncertain, self.mask_modules):
            # 1. multi-tea
            y_t = F.dropout(y_t, p=self.drop_rate, training=True)

            # 2. uncertainty weights
            y_s = Uncertain(y_s)
            y_t = Uncertain(y_t)

            # predict the masks
            mask = mask_module.forward_mask(y_t)
            if self.custom_mask and self._iter >= self.custom_mask_warmup:
                if self._iter == self.custom_mask_warmup:
                    print('Start customizing masks using student\'s masks.')
                with torch.no_grad():
                    mask_s = mask_module.forward_mask(y_s)  # [N, T, H, W]
                mask = mask * mask_s

            # get the masked features
            masked_y_s = y_s.unsqueeze(1) * \
                mask.unsqueeze(2)  # [N, n_masks, C, H, W]
            masked_y_t = y_t.unsqueeze(1) * \
                mask.unsqueeze(2)  # [N, n_masks, C, H, W]

            # masked distillation
            loss = (masked_y_s - masked_y_t)**2
            loss = loss.sum((3, 4))  # [N, n_masks, C]
            loss = loss / mask.sum((2, 3)).unsqueeze(-1)
            if self.weight_mask:
                weights = mask_module.forward_prob(y_t).flatten(1)  # [N, T]
                loss = (loss.mean(-1) * weights).sum(-1)
            loss = loss.mean()
            losses.append(loss)

        loss = sum(losses)
        self._iter += 1
        return self.loss_weight * loss
mmrazor/models/losses/akd_maskd.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
import torch.nn.functional as F

from ..builder import LOSSES
import math

class MaskModule(nn.Module):
    def __init__(self, channels):
        super(MaskModule, self).__init__()
        self.generation = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True), 
            nn.Conv2d(channels, channels, kernel_size=3, padding=1))

    def forward(self, x):

        return self.generation(x)

@LOSSES.register_module()
class AKD_MGDLoss(nn.Module):
    def __init__(self, alpha_mgd=0.00002, lambda_mgd=0.65, teacher_channels=[2048], drop_rate=0.1):
        super(AKD_MGDLoss, self).__init__()

        self.mask_modules = nn.ModuleList([
            MaskModule(channels=c) for c in teacher_channels]
        )
        self.alpha_mgd = alpha_mgd
        self.lambda_mgd = lambda_mgd

        self.Uncertain = nn.ModuleList([
            nn.BatchNorm2d(c, affine=False) for c in teacher_channels]
        )
        self.drop_rate = drop_rate

    def forward(self,
                y_s_list,
                y_t_list,):
        """Forward function.
        Args:
            preds_S(Tensor): Bs*C*H*W, student's feature map
            preds_T(Tensor): Bs*C*H*W, teacher's feature map
        """
        if not isinstance(y_s_list, (tuple, list)):
            y_s_list = (y_s_list, )
            y_t_list = (y_t_list, )
        assert len(y_s_list) == len(y_t_list) == len(self.mask_modules)

        losses = []
        for idx, (preds_S, preds_T) in enumerate(zip(y_s_list, y_t_list)):
            assert preds_S.shape[-2:] == preds_T.shape[-2:]
            # 1. multi-tea
            preds_T = F.dropout(preds_T, p=self.drop_rate, training=True)

            # 2. uncertainty weights           
            preds_S = self.Uncertain[idx](preds_S)
            preds_T = self.Uncertain[idx](preds_T)

            loss = self.get_dis_loss(preds_S, preds_T, idx)
            losses.append(loss)

        loss = sum(losses)
        return loss * self.alpha_mgd

    def get_dis_loss(self, preds_S, preds_T, idx):
        loss_mse = nn.MSELoss(reduction='sum')
        N, C, H, W = preds_T.shape

        device = preds_S.device
        mat = torch.rand((N,1,H,W)).to(device)
        mat = torch.where(mat > 1-self.lambda_mgd, 0, 1).to(device)

        masked_fea = torch.mul(preds_S, mat)
        new_fea = self.mask_modules[idx](masked_fea)

        dis_loss = loss_mse(new_fea, preds_T)/N

        return dis_loss
mmrazor/models/losses/akd_mgd.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
import torch.nn.functional as F

from ..builder import LOSSES


@LOSSES.register_module()
class ChannelWiseDivergence(nn.Module):
    """PyTorch version of `Channel-wise Distillation for Semantic Segmentation.

    <https://arxiv.org/abs/2011.13256>`_.

    Args:
        tau (float): Temperature coefficient. Defaults to 1.0.
        loss_weight (float): Weight of loss. Defaults to 1.0.
    """

    def __init__(
        self,
        tau=1.0,
        loss_weight=1.0,
    ):
        super(ChannelWiseDivergence, self).__init__()
        self.tau = tau
        self.loss_weight = loss_weight

    def forward(self, preds_S, preds_T):
        """Forward computation.

        Args:
            preds_S (torch.Tensor): The student model prediction with
                shape (N, C, H, W).
            preds_T (torch.Tensor): The teacher model prediction with
                shape (N, C, H, W).

        Return:
            torch.Tensor: The calculated loss value.
        """
        assert preds_S.shape[-2:] == preds_T.shape[-2:]
        N, C, H, W = preds_S.shape

        softmax_pred_T = F.softmax(preds_T.view(-1, W * H) / self.tau, dim=1)

        logsoftmax = torch.nn.LogSoftmax(dim=1)
        loss = torch.sum(softmax_pred_T *
                         logsoftmax(preds_T.view(-1, W * H) / self.tau) -
                         softmax_pred_T *
                         logsoftmax(preds_S.view(-1, W * H) / self.tau)) * (
                             self.tau**2)

        loss = self.loss_weight * loss / (C * N)

        return loss

mmrazor/models/losses/cwd.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
import torch.nn.functional as F

from ..builder import LOSSES


@LOSSES.register_module()
class KLDivergence(nn.Module):
    """A measure of how one probability distribution Q is different from a
    second, reference probability distribution P.

    Args:
        tau (float): Temperature coefficient. Defaults to 1.0.
        reduction (str): Specifies the reduction to apply to the loss:
            ``'none'`` | ``'batchmean'`` | ``'sum'`` | ``'mean'``.
            ``'none'``: no reduction will be applied,
            ``'batchmean'``: the sum of the output will be divided by
                the batchsize,
            ``'sum'``: the output will be summed,
            ``'mean'``: the output will be divided by the number of
                elements in the output.
            Default: ``'batchmean'``
        loss_weight (float): Weight of loss. Defaults to 1.0.
    """

    def __init__(
        self,
        tau=1.0,
        reduction='batchmean',
        loss_weight=1.0,
    ):
        super(KLDivergence, self).__init__()
        self.tau = tau
        self.loss_weight = loss_weight

        accept_reduction = {'none', 'batchmean', 'sum', 'mean'}
        assert reduction in accept_reduction, \
            f'KLDivergence supports reduction {accept_reduction}, ' \
            f'but gets {reduction}.'
        self.reduction = reduction

    def forward(self, preds_S, preds_T):
        """Forward computation.

        Args:
            preds_S (torch.Tensor): The student model prediction with
                shape (N, C, H, W) or shape (N, C).
            preds_T (torch.Tensor): The teacher model prediction with
                shape (N, C, H, W) or shape (N, C).

        Return:
            torch.Tensor: The calculated loss value.
        """
        preds_T = preds_T.detach()
        softmax_pred_T = F.softmax(preds_T / self.tau, dim=1)
        logsoftmax_preds_S = F.log_softmax(preds_S / self.tau, dim=1)
        loss = (self.tau**2) * F.kl_div(
            logsoftmax_preds_S, softmax_pred_T, reduction=self.reduction)
        return self.loss_weight * loss

mmrazor/models/losses/kl_divergence.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
import torch.nn.functional as F

from ..builder import LOSSES


def euclidean_distance(pred, squared=False, eps=1e-12):
    """Calculate the Euclidean distance between the two examples in the output
    representation space.

    Args:
        pred (torch.Tensor): The prediction of the teacher or student with
            shape (N, C).
        squared (bool): Whether to calculate the squared Euclidean
            distance. Defaults to False.
        eps (float): The minimum Euclidean distance between the two
            examples. Defaults to 1e-12.
    """
    pred_square = pred.pow(2).sum(dim=-1)  # (N, )
    prod = torch.mm(pred, pred.t())  # (N, N)
    distance = (pred_square.unsqueeze(1) + pred_square.unsqueeze(0) -
                2 * prod).clamp(min=eps)  # (N, N)

    if not squared:
        distance = distance.sqrt()

    distance = distance.clone()
    distance[range(len(prod)), range(len(prod))] = 0
    return distance


def angle(pred):
    """Calculate the angle-wise relational potential which measures the angle
    formed by the three examples in the output representation space.

    Args:
        pred (torch.Tensor): The prediction of the teacher or student with
            shape (N, C).
    """
    pred_vec = pred.unsqueeze(0) - pred.unsqueeze(1)  # (N, N, C)
    norm_pred_vec = F.normalize(pred_vec, p=2, dim=2)
    angle = torch.bmm(norm_pred_vec,
                      norm_pred_vec.transpose(1, 2)).view(-1)  # (N*N*N, )
    return angle


@LOSSES.register_module()
class DistanceWiseRKD(nn.Module):
    """PyTorch version of distance-wise loss of `Relational Knowledge
    Distillation.

    <https://arxiv.org/abs/1904.05068>`_.

    Args:
        loss_weight (float): Weight of distance-wise distillation loss.
            Defaults to 25.0.
        with_l2_norm (bool): Whether to normalize the model predictions before
            calculating the loss. Defaults to True.
    """

    def __init__(self, loss_weight=25.0, with_l2_norm=True):
        super(DistanceWiseRKD, self).__init__()

        self.loss_weight = loss_weight
        self.with_l2_norm = with_l2_norm

    def distance_loss(self, preds_S, preds_T):
        """Calculate distance-wise distillation loss."""
        d_T = euclidean_distance(preds_T, squared=False)
        # mean_d_T is a normalization factor for distance
        mean_d_T = d_T[d_T > 0].mean()
        d_T = d_T / mean_d_T

        d_S = euclidean_distance(preds_S, squared=False)
        mean_d_S = d_S[d_S > 0].mean()
        d_S = d_S / mean_d_S

        return F.smooth_l1_loss(d_S, d_T)

    def forward(self, preds_S, preds_T):
        """Forward computation.

        Args:
            preds_S (torch.Tensor): The student model prediction with
                shape (N, C, H, W) or shape (N, C).
            preds_T (torch.Tensor): The teacher model prediction with
                shape (N, C, H, W) or shape (N, C).
        Return:
            torch.Tensor: The calculated loss value.
        """
        preds_S = preds_S.view(preds_S.shape[0], -1)
        preds_T = preds_T.view(preds_T.shape[0], -1)
        if self.with_l2_norm:
            preds_S = F.normalize(preds_S, p=2, dim=1)
            preds_T = F.normalize(preds_T, p=2, dim=1)

        loss = self.distance_loss(preds_S, preds_T) * self.loss_weight

        return loss


@LOSSES.register_module()
class AngleWiseRKD(nn.Module):
    """PyTorch version of angle-wise loss of `Relational Knowledge
    Distillation.

    <https://arxiv.org/abs/1904.05068>`_.

    Args:
        loss_weight (float): Weight of angle-wise distillation loss.
            Defaults to 50.0.
        with_l2_norm (bool): Whether to normalize the model predictions before
            calculating the loss. Defaults to True.
    """

    def __init__(self, loss_weight=50.0, with_l2_norm=True):
        super(AngleWiseRKD, self).__init__()

        self.loss_weight = loss_weight
        self.with_l2_norm = with_l2_norm

    def angle_loss(self, preds_S, preds_T):
        """Calculate the angle-wise distillation loss."""
        angle_T = angle(preds_T)
        angle_S = angle(preds_S)
        return F.smooth_l1_loss(angle_S, angle_T)

    def forward(self, preds_S, preds_T):
        """Forward computation.

        Args:
            preds_S (torch.Tensor): The student model prediction with
                shape (N, C, H, W) or shape (N, C).
            preds_T (torch.Tensor): The teacher model prediction with
                shape (N, C, H, W) or shape (N, C).
        Return:
            torch.Tensor: The calculated loss value.
        """
        preds_S = preds_S.view(preds_S.shape[0], -1)
        preds_T = preds_T.view(preds_T.shape[0], -1)
        if self.with_l2_norm:
            preds_S = F.normalize(preds_S, p=2, dim=-1)
            preds_T = F.normalize(preds_T, p=2, dim=-1)

        loss = self.angle_loss(preds_S, preds_T) * self.loss_weight

        return loss

mmrazor/models/losses/relational_kd.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
import torch.nn.functional as F

from ..builder import LOSSES


@LOSSES.register_module()
class WSLD(nn.Module):
    """PyTorch version of `Rethinking Soft Labels for Knowledge
    Distillation: A Bias-Variance Tradeoff Perspective
    <https://arxiv.org/abs/2102.00650>`_.

    Args:
        tau (float): Temperature coefficient. Defaults to 1.0.
        loss_weight (float): Weight of loss. Defaults to 1.0.
        num_classes (int): Defaults to 1000.
    """

    def __init__(self, tau=1.0, loss_weight=1.0, num_classes=1000):
        super(WSLD, self).__init__()

        self.tau = tau
        self.loss_weight = loss_weight
        self.num_classes = num_classes
        self.softmax = nn.Softmax(dim=1).cuda()
        self.logsoftmax = nn.LogSoftmax(dim=1).cuda()

    def forward(self, student, teacher):

        gt_labels = self.current_data['gt_label']

        student_logits = student / self.tau
        teacher_logits = teacher / self.tau

        teacher_probs = self.softmax(teacher_logits)

        ce_loss = -torch.sum(
            teacher_probs * self.logsoftmax(student_logits), 1, keepdim=True)

        student_detach = student.detach()
        teacher_detach = teacher.detach()
        log_softmax_s = self.logsoftmax(student_detach)
        log_softmax_t = self.logsoftmax(teacher_detach)
        one_hot_labels = F.one_hot(
            gt_labels, num_classes=self.num_classes).float()
        ce_loss_s = -torch.sum(one_hot_labels * log_softmax_s, 1, keepdim=True)
        ce_loss_t = -torch.sum(one_hot_labels * log_softmax_t, 1, keepdim=True)

        focal_weight = ce_loss_s / (ce_loss_t + 1e-7)
        ratio_lower = torch.zeros(1).cuda()
        focal_weight = torch.max(focal_weight, ratio_lower)
        focal_weight = 1 - torch.exp(-focal_weight)
        ce_loss = focal_weight * ce_loss

        loss = (self.tau**2) * torch.mean(ce_loss)

        loss = self.loss_weight * loss

        return loss

mmrazor/models/losses/weighted_soft_label_distillation.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .mutable_edge import DifferentiableEdge, GumbelEdge, MutableEdge
from .mutable_module import MutableModule
from .mutable_op import DifferentiableOP, GumbelOP, MutableOP, OneShotOP

__all__ = [
    'MutableModule', 'MutableOP', 'MutableEdge', 'DifferentiableOP',
    'DifferentiableEdge', 'GumbelEdge', 'GumbelOP', 'OneShotOP'
]

mmrazor/models/mutables/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy

import torch
import torch.nn as nn
import torch.nn.functional as F

from ..builder import MUTABLES
from .mutable_module import MutableModule


class MutableEdge(MutableModule):
    """Mutable Edge. In some NAS algorithms (Darts, AutoDeeplab, etc.), the
    connections between modules are searchable, such as the connections between
    a node and its previous nodes in Darts. ``MutableEdge`` has N modules to
    process N inputs respectively.

    Args:
        choices (torch.nn.ModuleDict): Unlike ``MutableOP``, there are already
            created modules in choices.
    """

    def __init__(self, choices, **kwargs):
        super(MutableEdge, self).__init__(**kwargs)
        assert isinstance(choices, nn.ModuleDict), \
            f'In {self.space_id}, the choices must be torch.nn.ModuleDict.'
        self.choices = choices
        self.choice_mask = self.build_choice_mask()
        # Once execute ``export``, the unchosen modules will be removed.
        # Record full choice names, which will be used in forward.
        self.full_choice_names = copy.deepcopy(self.choice_names)

    def build_choices(self, cfg):
        """MutableEdge's choices is already built."""
        pass


@MUTABLES.register_module()
class DifferentiableEdge(MutableEdge):
    """Differentiable Edge.

    Search the best module from choices by learnable parameters.

    Args:
        with_arch_param (bool): whether build learable architecture parameters.
    """

    def __init__(self, with_arch_param, **kwargs):
        super(DifferentiableEdge, self).__init__(**kwargs)
        self.with_arch_param = with_arch_param

    def build_arch_param(self):
        """build learnable architecture parameters."""
        if self.with_arch_param:
            return nn.Parameter(torch.randn(self.num_choices) * 1e-3)
        else:
            return None

    def compute_arch_probs(self, arch_param):
        """compute chosen probs according architecture parameters."""
        return F.softmax(arch_param, -1)

    def forward(self, prev_inputs, arch_param=None):
        """forward function.

        In some algorithms, there are several ``MutableModule`` share the same
        architecture parameters. So the architecture parameters are passed
        in as args.

        Args:
            prev_inputs (list[torch.Tensor]): each choice's inputs.
            arch_param (torch.nn.Parameter): architecture parameters.
        """

        assert len(self.full_choice_names) == len(prev_inputs), \
            f'In {self.space_id}, the length of full choice names must be \
                same as the length of previous inputs.'

        if self.with_arch_param:
            assert arch_param is not None, \
                f'In {self.space_id}, the arch_param can not be None when the \
                    with_arch_param=True.'

            # 1. compute choices' probs.
            probs = self.compute_arch_probs(arch_param)

            # 2. process each input
            outputs = list()
            for prob, module, input in zip(probs, self.choice_modules,
                                           prev_inputs):
                if prob > 0:
                    outputs.append(prob * module(input))

        else:
            outputs = list()
            for name, input in zip(self.full_choice_names, prev_inputs):
                if name not in self.choice_names:
                    continue
                module = self.choices[name]
                outputs.append(module(input))

            assert len(outputs) > 0

        return sum(outputs)


@MUTABLES.register_module()
class GumbelEdge(DifferentiableEdge):
    """Gumbel Edge.

    Search the best module from choices by gumbel trick.
    """

    def __init__(self, **kwargs):
        super(GumbelEdge, self).__init__(**kwargs)

    def set_temperature(self, temperature):
        """Modify the temperature."""
        self.temperature = temperature

    def compute_arch_probs(self, arch_param):
        """compute chosen probs by gumbel trick."""
        probs = F.gumbel_softmax(
            arch_param, tau=self.tau, hard=self.hard, dim=-1)
        return probs

mmrazor/models/mutables/mutable_edge.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod

import torch
import torch.nn as nn
from mmcv.runner import BaseModule


class MutableModule(BaseModule, metaclass=ABCMeta):
    """Base class for ``MUTABLES``. Searchable module for building searchable
    architecture in NAS. It mainly consists of module and mask, and achieving
    searchable function by handling mask.

    Args:
        space_id (str): Used to index ``Placeholder``, it is one and only index
            for each ``Placeholder``.
        num_chosen (str): The number of chosen ``OPS`` in the ``MUTABLES``.
        init_cfg (dict): Init config for ``BaseModule``.
    """

    def __init__(self, space_id, num_chosen=1, init_cfg=None, **kwargs):
        super(MutableModule, self).__init__(init_cfg)
        self.space_id = space_id
        self.num_chosen = num_chosen

    @abstractmethod
    def forward(self, x):
        """Forward computation.

        Args:
            x (tensor | tuple[tensor]): x could be a Torch.tensor or a tuple of
                Torch.tensor, containing input data for forward computation.
        """
        pass

    @abstractmethod
    def build_choices(self, cfg):
        """Build all chosen ``OPS`` used to combine ``MUTABLES``, and the
        choices will be sampled.

        Args:
            cfg (dict): The config for the choices.
        """
        pass

    def build_choice_mask(self):
        """Generate the choice mask for the choices of ``MUTABLES``.

        Returns:
            torch.Tensor: Init choice mask. Its elements' type is bool.
        """
        if torch.cuda.is_available():
            return torch.ones(self.num_choices).bool().cuda()
        else:
            return torch.ones(self.num_choices).bool()

    def set_choice_mask(self, mask):
        """Use the mask to update the choice mask.

        Args:
            mask (torch.Tensor): Choice mask specified to update the choice
                mask.
        """
        assert self.choice_mask.size(0) == mask.size(0)
        self.choice_mask = mask

    @property
    def num_choices(self):
        """The number of the choices.

        Returns:
            int: the length of the choices.
        """
        return len(self.choices)

    @property
    def choice_names(self):
        """The choices' names.

        Returns:
            tuple: The keys of the choices.
        """
        assert isinstance(self.choices, nn.ModuleDict), \
            'candidates must be nn.ModuleDict.'
        return tuple(self.choices.keys())

    @property
    def choice_modules(self):
        """The choices' modules.

        Returns:
            tuple: The values of the choices.
        """
        assert isinstance(self.choices, nn.ModuleDict), \
            'candidates must be nn.ModuleDict.'
        return tuple(self.choices.values())

    def build_space_mask(self):
        """Generate the space mask for the search spaces of ``MUTATORS``.

        Returns:
            torch.Tensor: Init choice mask. Its elements' type is float.
        """
        if torch.cuda.is_available():
            return torch.ones(self.num_choices).cuda() * 1.0
        else:
            return torch.ones(self.num_choices) * 1.0

    def export(self, chosen):
        """Delete not chosen ``OPS`` in the choices.

        Args:
            chosen (list[str]): Names of chosen ``OPS``.
        """
        for name in self.choice_names:
            if name not in chosen:
                self.choices.pop(name)

mmrazor/models/mutables/mutable_module.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy

import torch
import torch.nn as nn
import torch.nn.functional as F

from ..builder import MUTABLES, build_op
from .mutable_module import MutableModule


class MutableOP(MutableModule):
    """An important type of ``MUTABLES``, inherits from ``MutableModule``.

    Args:
        choices (dict): The configs for the choices, the chosen ``OPS`` used to
            combine ``MUTABLES``.
        choice_args (dict): The args used to set chosen ``OPS``.
    """

    def __init__(self, choices, choice_args, **kwargs):
        super(MutableOP, self).__init__(**kwargs)
        self.choices = self.build_choices(choices, choice_args)
        self.choice_mask = self.build_choice_mask()
        self.full_choice_names = copy.deepcopy(self.choice_names)

    def build_choices(self, cfgs, choice_args):
        """Build all chosen ``OPS`` used to combine ``MUTABLES``, and the
        choices will be sampled.

        Args:
            cfgs (dict): The configs for the choices.
            choice_args (dict): The args used to set chosen ``OPS``.

        Returns:
            torch.nn.ModuleDict: Consists of chosen ``OPS`` in the arg `cfgs`.
        """
        choices = nn.ModuleDict()
        for name, cfg in cfgs.items():
            cfg.update(choice_args)
            op_module = build_op(cfg)
            choices.add_module(name, op_module)
        return choices


@MUTABLES.register_module()
class OneShotOP(MutableOP):
    """A type of ``MUTABLES`` for the one-shot NAS."""

    def __init__(self, **kwargs):
        super(OneShotOP, self).__init__(**kwargs)
        assert self.num_chosen == 1

    def forward(self, x):
        """Forward computation for chosen ``OPS``, in one-shot NAS, the number
        of chosen ``OPS`` can only be one.

        Args:
            x (tensor | tuple[tensor]): x could be a Torch.tensor or a tuple of
                Torch.tensor, containing input data for forward computation.

        Returns:
            torch.Tensor: The result of forward.
        """
        outputs = list()
        for name, chosen_bool in zip(self.full_choice_names, self.choice_mask):
            if name not in self.choice_names:
                continue
            if not chosen_bool:
                continue
            module = self.choices[name]
            outputs.append(module(x))

        assert len(outputs) > 0

        return sum(outputs)


@MUTABLES.register_module()
class DifferentiableOP(MutableOP):
    """Differentiable OP.

    Search the best module from choices by learnable parameters.

    Args:
        with_arch_param (bool): whether build learable architecture parameters.
    """

    def __init__(self, with_arch_param, **kwargs):
        super(DifferentiableOP, self).__init__(**kwargs)
        self.with_arch_param = with_arch_param

    def build_arch_param(self):
        """build learnable architecture parameters."""
        if self.with_arch_param:
            return nn.Parameter(torch.randn(self.num_choices) * 1e-3)
        else:
            return None

    def compute_arch_probs(self, arch_param):
        """compute chosen probs according architecture parameters."""
        return F.softmax(arch_param, -1)

    def forward(self, x, arch_param=None):
        """forward function.

        In some algorithms, there are several ``MutableModule`` share the same
        architecture parameters. So the architecture parameters are passed
        in as args.

        Args:
            prev_inputs (list[torch.Tensor]): each choice's inputs.
            arch_param (torch.nn.Parameter): architecture parameters.
        """
        if self.with_arch_param:
            assert arch_param is not None, \
                f'In {self.space_id}, the arch_param can not be None when the \
                    with_arch_param=True.'

            # 1. compute choices' probs.
            probs = self.compute_arch_probs(arch_param)

            # 2. compute every op's outputs.
            outputs = list()
            for prob, module in zip(probs, self.choice_modules):
                if prob > 0:
                    outputs.append(prob * module(x))

        else:
            outputs = list()
            for name, chosen_bool in zip(self.full_choice_names,
                                         self.choice_mask):
                if name not in self.choice_names:
                    continue
                if not chosen_bool:
                    continue
                module = self.choices[name]
                outputs.append(module(x))

            assert len(outputs) > 0
        return sum(outputs)


@MUTABLES.register_module()
class GumbelOP(DifferentiableOP):
    """Gumbel OP.

    Search the best module from choices by gumbel trick.
    """

    def __init__(self, tau=1.0, hard=True, **kwargs):
        super(GumbelOP, self).__init__(**kwargs)
        self.tau = tau
        self.hard = hard

    def set_temperature(self, temperature):
        """Modify the temperature."""
        self.temperature = temperature

    def compute_arch_probs(self, arch_param):
        """compute chosen probs by gumbel trick."""
        probs = F.gumbel_softmax(
            arch_param, tau=self.tau, hard=self.hard, dim=-1)
        return probs

mmrazor/models/mutables/mutable_op.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .darts_mutator import DartsMutator
from .differentiable_mutator import DifferentiableMutator
from .one_shot_mutator import OneShotMutator

__all__ = ['DifferentiableMutator', 'DartsMutator', 'OneShotMutator']

mmrazor/models/mutators/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta

from mmcv.runner import BaseModule

from mmrazor.models.architectures import Placeholder
from mmrazor.models.builder import MUTABLES, MUTATORS
from mmrazor.models.mutables import MutableModule


@MUTATORS.register_module()
class BaseMutator(BaseModule, metaclass=ABCMeta):
    """Base class for mutators."""

    def __init__(self, placeholder_mapping=None, init_cfg=None):
        super(BaseMutator, self).__init__(init_cfg=init_cfg)
        self.placeholder_mapping = placeholder_mapping

    def prepare_from_supernet(self, supernet):
        """Implement some preparatory work based on supernet, including
        ``convert_placeholder`` and ``build_search_spaces``.

        Args:
            supernet (:obj:`torch.nn.Module`): The architecture to be used
                in your algorithm.
        """
        if self.placeholder_mapping is not None:
            self.convert_placeholder(supernet, self.placeholder_mapping)

        self.search_spaces = self.build_search_spaces(supernet)

    def build_search_spaces(self, supernet):
        """Build a search space from the supernet.

        Args:
            supernet (:obj:`torch.nn.Module`): The architecture to be used
                in your algorithm.

        Returns:
            dict: To collect some information about ``MutableModule`` in the
                supernet.
        """
        search_spaces = dict()

        def traverse(module):
            for child in module.children():
                if isinstance(child, MutableModule):
                    if child.space_id not in search_spaces.keys():
                        search_spaces[child.space_id] = dict(
                            modules=[child],
                            choice_names=child.choice_names,
                            num_chosen=child.num_chosen,
                            space_mask=child.build_space_mask())
                    else:
                        search_spaces[child.space_id]['modules'].append(child)

                traverse(child)

        traverse(supernet)
        return search_spaces

    def convert_placeholder(self, supernet, placeholder_mapping):
        """Replace all placeholders in the model.

        Args:
            supernet (:obj:`torch.nn.Module`): The architecture to be used in
                your algorithm.
            placeholder_mapping (dict): Record which placeholders need to be
                replaced by which ops,
                its keys are the properties ``placeholder_group`` of
                placeholders used in the searchable architecture,
                its values are the registered ``OPS``.
        """

        def traverse(module):

            for name, child in module.named_children():
                if isinstance(child, Placeholder):

                    mutable_cfg = placeholder_mapping[
                        child.placeholder_group].copy()
                    assert 'type' in mutable_cfg, f'{mutable_cfg}'
                    mutable_type = mutable_cfg.pop('type')
                    assert mutable_type in MUTABLES, \
                        f'{mutable_type} not in MUTABLES.'
                    mutable_constructor = MUTABLES.get(mutable_type)
                    mutable_kwargs = child.placeholder_kwargs
                    mutable_kwargs.update(mutable_cfg)
                    mutable_module = mutable_constructor(**mutable_kwargs)
                    setattr(module, name, mutable_module)

                    # setattr(module, name, choice_module)
                    # If the new MUTABLE is MutableEdge, it may have MutableOP,
                    # so here we need to traverse the new MUTABLES.
                    traverse(mutable_module)
                else:
                    traverse(child)

        traverse(supernet)

    def deploy_subnet(self, supernet, subnet_dict):
        """Export the subnet from the supernet based on the specified
        subnet_dict.

        Args:
            supernet (:obj:`torch.nn.Module`): The architecture to be used in
                your algorithm.
            subnet_dict (dict): Record the information to build the subnet from
                the supernet,
                its keys are the properties ``space_id`` of placeholders in the
                mutator's search spaces,
                its values are dicts: {'chosen': ['chosen name1',
                'chosen name2', ...]}
        """

        def traverse(module):
            for name, child in module.named_children():
                if isinstance(child, MutableModule):
                    space_id = child.space_id
                    chosen = subnet_dict[space_id]['chosen']
                    child.export(chosen)

                traverse(child)

        traverse(supernet)

mmrazor/models/mutators/base.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.

import torch
from torch.nn import functional as F

from mmrazor.models.builder import MUTATORS
from .differentiable_mutator import DifferentiableMutator


@MUTATORS.register_module()
class DartsMutator(DifferentiableMutator):

    def __init__(self, ignore_choices=('zero', ), **kwargs):
        super().__init__(**kwargs)
        self.ignore_choices = ignore_choices

    def search_subnet(self):

        subnet_dict = dict()
        for space_id, sub_space in self.search_spaces.items():
            if space_id in self.arch_params:
                space_arch_param = self.arch_params[space_id]
                arch_probs = F.softmax(space_arch_param, dim=-1)
                choice_names = sub_space['choice_names']
                keep_idx = [
                    i for i, name in enumerate(choice_names)
                    if name not in self.ignore_choices
                ]
                best_choice_prob, best_choice_idx = torch.max(
                    arch_probs[keep_idx], 0)
                best_choice_idx = keep_idx[best_choice_idx.item()]
                best_choice_name = choice_names[best_choice_idx]

                subnet_dict[space_id] = dict(
                    chosen=[best_choice_name],
                    chosen_probs=[best_choice_prob.item()])

        def sort_key(x):
            return subnet_dict[x]['chosen_probs'][0]

        for space_id, sub_space in self.search_spaces.items():
            if space_id not in self.arch_params:
                num_chosen = sub_space['num_chosen']
                choice_names = sub_space['choice_names']
                sorted_edges = list(
                    sorted(choice_names, key=sort_key, reverse=True))
                chosen = sorted_edges[:num_chosen]
                subnet_dict[space_id] = dict(chosen=chosen)

                for not_chosen in sorted_edges[num_chosen:]:
                    subnet_dict.pop(not_chosen)

        return subnet_dict

mmrazor/models/mutators/darts_mutator.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from abc import abstractmethod
from functools import partial

from torch import nn

from mmrazor.models.builder import MUTATORS
from mmrazor.models.mutables import MutableModule
from .base import BaseMutator


@MUTATORS.register_module()
class DifferentiableMutator(BaseMutator):
    """A mutator for the differentiable NAS, which mainly provide some core
    functions of changing the structure of ``ARCHITECTURES``."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def prepare_from_supernet(self, supernet):
        """Inherit from ``BaseMutator``'s, execute some customized functions
        exclude implementing origin ``prepare_from_supernet``.

        Args:
            supernet (:obj:`torch.nn.Module`): The architecture to be used
                in your algorithm.
        """
        super().prepare_from_supernet(supernet)
        self.arch_params = self.build_arch_params(supernet)
        self.modify_supernet_forward(supernet)

    def build_arch_params(self, supernet):
        """This function will build many arch params, which are generally used
        in diffirentiale search algorithms, such as Darts' series. Each
        space_id corresponds to an arch param, so the Mutable with the same
        space_id share the same arch param.

        Args:
            supernet (:obj:`torch.nn.Module`): The architecture to be used
                in your algorithm.

        Returns:
            torch.nn.ParameterDict: the arch params are got after traversing
                the supernet.
        """

        arch_params = nn.ParameterDict()

        # Traverse all the child modules of the model. If a child module is an
        # Space instance and its space_id is not recorded, call its
        # :func:'build_space_architecture' and record the return value. If not,
        # pass.
        def traverse(module):
            for name, child in module.named_children():

                if isinstance(child, MutableModule):
                    space_id = child.space_id
                    if space_id not in arch_params:
                        space_arch_param = child.build_arch_param()
                        if space_arch_param is not None:
                            arch_params[space_id] = space_arch_param

                traverse(child)

        traverse(supernet)

        return arch_params

    def modify_supernet_forward(self, supernet):
        """Modify the supernet's default value in forward. Traverse all child
        modules of the model, modify the supernet's default value in
        :func:'forward' of each Space.

        Args:
            supernet (:obj:`torch.nn.Module`): The architecture to be used
                in your algorithm.
        """

        def traverse(module):
            for name, child in module.named_children():

                if isinstance(child, MutableModule):
                    if child.space_id in self.arch_params.keys():
                        space_id = child.space_id
                        space_arch_param = self.arch_params[space_id]
                        child.forward = partial(
                            child.forward, arch_param=space_arch_param)

                traverse(child)

        traverse(supernet)

    @abstractmethod
    def search_subnet(self):
        pass

mmrazor/models/mutators/differentiable_mutator.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from functools import partial

import numpy as np
import torch
import torch.distributed as dist

from mmrazor.models.builder import MUTATORS
from .base import BaseMutator


@MUTATORS.register_module()
class OneShotMutator(BaseMutator):
    """A mutator for the one-shot NAS, which mainly provide some core functions
    of changing the structure of ``ARCHITECTURES``."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @staticmethod
    def get_random_mask(space_info, searching):
        """Generate random mask for randomly sampling.

        Args:
            space_info (dict): Record the information of the space need
                to sample.
            searching (bool): Whether is in search stage.

        Returns:
            torch.Tensor: Random mask generated.
        """
        space_mask = space_info['space_mask']
        num_chosen = space_info['num_chosen']
        assert num_chosen <= space_mask.size()[0]
        choice_idx = torch.multinomial(space_mask, num_chosen)
        choice_mask = torch.zeros_like(space_mask)
        choice_mask[choice_idx] = 1
        if dist.is_available() and dist.is_initialized() and not searching:
            dist.broadcast(choice_mask, src=0)
        return choice_mask

    def sample_subnet(self, searching=False):
        """Random sample subnet by random mask.

        Args:
            searching (bool): Whether is in search stage.

        Returns:
            dict: Record the information to build the subnet from the supernet,
                its keys are the properties ``space_id`` of placeholders in the
                mutator's search spaces,
                its values are random mask generated.
        """
        subnet_dict = dict()
        for space_id, space_info in self.search_spaces.items():
            subnet_dict[space_id] = self.get_random_mask(space_info, searching)
        return subnet_dict

    def set_subnet(self, subnet_dict):
        """Setting subnet in the supernet based on the result of
        ``sample_subnet`` by changing the flag: ``in_subnet``, which is easy to
        implement some operations for subnet, such as ``forward``, calculate
        flops and so on.

        Args:
            subnet_dict (dict): Record the information to build the subnet
                from the supernet,
                its keys are the properties ``space_id`` of placeholders in the
                mutator's search spaces,
                its values are masks.
        """
        for space_id, space_info in self.search_spaces.items():
            choice_mask = subnet_dict[space_id]
            for module in space_info['modules']:
                module.choice_mask = choice_mask
                for i, choice in enumerate(module.choices.values()):
                    if choice_mask[i]:
                        choice.apply(
                            partial(self.reset_in_subnet, in_subnet=True))
                    else:
                        choice.apply(
                            partial(self.reset_in_subnet, in_subnet=False))

    @staticmethod
    def reset_in_subnet(m, in_subnet=True):
        """Reset the module's attribution.

        Args:
            m (:obj:`torch.nn.Module`): The module in the supernet.
            in_subnet (bool): If the module in subnet, set ``in_subnet`` to
                True, otherwise set to False.
        """
        m.__in_subnet__ = in_subnet

    def set_chosen_subnet(self, subnet_dict):
        """Set chosen subnet in the search_spaces after searching stage.

        Args:
            subnet_dict (dict): Record the information to build the subnet from
                the supernet,
                its keys are the properties ``space_id`` of placeholders in the
                mutator's search spaces,
                its values are masks.
        """
        for space_id, mask in subnet_dict.items():
            idxs = [i for i, x in enumerate(mask.tolist()) if x == 1.0]
            self.search_spaces[space_id]['chosen'] = [
                self.search_spaces[space_id]['choice_names'][i] for i in idxs
            ]

    def mutation(self, subnet_dict, prob=0.1):
        """Mutation used in evolution search.

        Args:
            subnet_dict (dict): Record the information to build the subnet
                from the supernet, its keys are the properties ``space_id``
                of placeholders in the mutator's search spaces, its values
                are masks.
            prob (float): The probability of mutation.

        Returns:
            dict: A new subnet_dict after mutation.
        """
        mutation_subnet_dict = copy.deepcopy(subnet_dict)
        for name, mask in subnet_dict.items():
            if np.random.random_sample() < prob:
                mutation_subnet_dict[name] = self.get_random_mask(
                    self.search_spaces[name], searching=True)
        return mutation_subnet_dict

    @staticmethod
    def crossover(subnet_dict1, subnet_dict2):
        """Crossover used in evolution search.

        Args:
            subnet_dict1 (dict): Record the information to build the subnet
                from the supernet,
                its keys are the properties ``space_id`` of placeholders in the
                mutator's search spaces,
                its values are masks.
            subnet_dict2 (dict): Record the information to build the subnet
                from the supernet,
                its keys are the properties ``space_id`` of placeholders in the
                mutator's search spaces,
                its values are masks.

        Returns:
            dict: A new subnet_dict after crossover.
        """
        crossover_subnet_dict = copy.deepcopy(subnet_dict1)
        for name, mask in subnet_dict2.items():
            if np.random.random_sample() < 0.5:
                crossover_subnet_dict[name] = mask
        return crossover_subnet_dict

mmrazor/models/mutators/one_shot_mutator.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .common import Identity
from .darts_series import (DartsDilConv, DartsPoolBN, DartsSepConv,
                           DartsSkipConnect, DartsZero)
from .mobilenet_series import MBBlock
from .shufflenet_series import ShuffleBlock, ShuffleXception

__all__ = [
    'ShuffleBlock', 'ShuffleXception', 'DartsPoolBN', 'DartsDilConv',
    'DartsSepConv', 'DartsSkipConnect', 'DartsZero', 'MBBlock', 'Identity'
]

mmrazor/models/ops/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.runner import BaseModule


class BaseOP(BaseModule):
    """Base class for searchable operations.

    Args:
        in_channels (int): The input channels of the operation.
        out_channels (int): The output channels of the operation.
        stride (int): Stride of the operation. Defaults to 1.
    """

    def __init__(self, in_channels, out_channels, stride=1, **kwargs):
        super(BaseOP, self).__init__(**kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.stride = stride

mmrazor/models/ops/base.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.cnn import ConvModule

from ..builder import OPS
from .base import BaseOP


@OPS.register_module()
class Identity(BaseOP):
    """Base class for searchable operations.

    Args:
        conv_cfg (dict, optional): Config dict for convolution layer.
            Default: None, which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='BN').
        act_cfg (dict): Config dict for activation layer.
            Default: None.
    """

    def __init__(self,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=None,
                 **kwargs):
        super(Identity, self).__init__(**kwargs)

        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg

        if self.stride != 1 or self.in_channels != self.out_channels:
            self.downsample = ConvModule(
                self.in_channels,
                self.out_channels,
                kernel_size=1,
                stride=self.stride,
                padding=0,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg)
        else:
            self.downsample = None

    def forward(self, x):
        if self.downsample is not None:
            x = self.downsample(x)
        return x

mmrazor/models/ops/common.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from mmcv.cnn import build_norm_layer
from mmcv.cnn.bricks import DropPath

from ..builder import OPS
from .base import BaseOP


@OPS.register_module()
class DartsPoolBN(BaseOP):

    def __init__(self,
                 pool_type,
                 kernel_size=3,
                 norm_cfg=dict(type='BN'),
                 use_drop_path=False,
                 **kwargs):
        super(DartsPoolBN, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.norm_cfg = norm_cfg
        if pool_type == 'max':
            self.pool = nn.MaxPool2d(self.kernel_size, self.stride, 1)
        elif pool_type == 'avg':
            self.pool = nn.AvgPool2d(
                self.kernel_size, self.stride, 1, count_include_pad=False)
        self.bn = build_norm_layer(self.norm_cfg, self.out_channels)[1]

        if use_drop_path:
            self.drop_path = DropPath()
        else:
            self.drop_path = None

    def forward(self, x):
        out = self.pool(x)
        out = self.bn(out)
        if self.drop_path is not None:
            out = self.drop_path(out)

        return out


@OPS.register_module()
class DartsDilConv(BaseOP):

    def __init__(self,
                 kernel_size,
                 use_drop_path=False,
                 norm_cfg=dict(type='BN'),
                 **kwargs):
        super(DartsDilConv, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.norm_cfg = norm_cfg
        self.dilation = 2
        assert self.kernel_size in [3, 5]
        assert self.stride in [1, 2]
        self.conv1 = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(
                self.in_channels,
                self.in_channels,
                self.kernel_size,
                self.stride, (self.kernel_size // 2) * self.dilation,
                dilation=self.dilation,
                groups=self.in_channels,
                bias=False),
            nn.Conv2d(
                self.in_channels, self.out_channels, 1, stride=1, bias=False),
            build_norm_layer(self.norm_cfg, self.in_channels)[1])

        if use_drop_path:
            self.drop_path = DropPath()
        else:
            self.drop_path = None

    def forward(self, x):
        out = self.conv1(x)
        if self.drop_path is not None:
            out = self.drop_path(out)
        return out


@OPS.register_module()
class DartsSepConv(BaseOP):

    def __init__(self,
                 kernel_size,
                 use_drop_path=False,
                 norm_cfg=dict(type='BN'),
                 **kwargs):
        super(DartsSepConv, self).__init__(**kwargs)

        self.kernel_size = kernel_size
        self.norm_cfg = norm_cfg
        assert self.kernel_size in [3, 5]
        assert self.stride in [1, 2]
        self.conv1 = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(
                self.in_channels,
                self.in_channels,
                self.kernel_size,
                self.stride,
                self.kernel_size // 2,
                groups=self.in_channels,
                bias=False),
            nn.Conv2d(
                self.in_channels, self.in_channels, 1, stride=1, bias=False),
            build_norm_layer(self.norm_cfg, self.in_channels)[1])
        self.conv2 = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(
                self.in_channels,
                self.out_channels,
                self.kernel_size,
                1,
                self.kernel_size // 2,
                groups=self.in_channels,
                bias=False),
            nn.Conv2d(
                self.out_channels, self.out_channels, 1, stride=1, bias=False),
            build_norm_layer(self.norm_cfg, self.out_channels)[1])

        if use_drop_path:
            self.drop_path = DropPath()
        else:
            self.drop_path = None

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        if self.drop_path is not None:
            out = self.drop_path(out)
        return out


@OPS.register_module()
class DartsSkipConnect(BaseOP):
    """Reduce feature map size by factorized pointwise (stride=2)."""

    def __init__(self,
                 use_drop_path=False,
                 norm_cfg=dict(type='BN'),
                 **kwargs):
        super(DartsSkipConnect, self).__init__(**kwargs)
        self.norm_cfg = norm_cfg
        if self.stride > 1:
            self.relu = nn.ReLU()
            self.conv1 = nn.Conv2d(
                self.in_channels,
                self.out_channels // 2,
                1,
                stride=2,
                padding=0,
                bias=False)
            self.conv2 = nn.Conv2d(
                self.in_channels,
                self.out_channels // 2,
                1,
                stride=2,
                padding=0,
                bias=False)
            self.bn = build_norm_layer(self.norm_cfg, self.out_channels)[1]

        if use_drop_path:
            self.drop_path = DropPath()
        else:
            self.drop_path = None

    def forward(self, x):
        if self.stride > 1:
            x = self.relu(x)
            out = torch.cat(
                [self.conv1(x), self.conv2(x[:, :, 1:, 1:])], dim=1)
            out = self.bn(out)
            if self.drop_path is not None:
                out = self.drop_path(out)
        else:
            out = x
        return out


@OPS.register_module()
class DartsZero(BaseOP):

    def __init__(self, **kwargs):
        super(DartsZero, self).__init__(**kwargs)

    def forward(self, x):
        if self.stride == 1:
            return x.mul(0.)
        return x[:, :, ::self.stride, ::self.stride].mul(0.)

mmrazor/models/ops/darts_series.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
import torch.utils.checkpoint as cp
from mmcls.models.utils import SELayer
from mmcv.cnn import ConvModule
from mmcv.cnn.bricks import DropPath

from ..builder import OPS
from .base import BaseOP


@OPS.register_module()
class MBBlock(BaseOP):
    """Mobilenet block for Searchable backbone.

    Args:
        kernel_size (int): Size of the convolving kernel.
        expand_ratio (int): The input channels' expand factor of the depthwise
             convolution.
        se_cfg (dict, optional): Config dict for se layer. Defaults to None,
            which means no se layer.
        conv_cfg (dict, optional): Config dict for convolution layer.
            Default: None, which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='BN').
        act_cfg (dict): Config dict for activation layer.
            Default: dict(type='ReLU').
        drop_path_rate (float): stochastic depth rate. Defaults to 0.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.

    Returns:
        Tensor: The output tensor.
    """

    def __init__(self,
                 kernel_size,
                 expand_ratio,
                 se_cfg=None,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 drop_path_rate=0.,
                 with_cp=False,
                 **kwargs):

        super(MBBlock, self).__init__(**kwargs)
        self.with_res_shortcut = (
            self.stride == 1 and self.in_channels == self.out_channels)
        assert self.stride in [1, 2]
        self.kernel_size = kernel_size
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.with_cp = with_cp
        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate > 0 else nn.Identity()
        self.with_se = se_cfg is not None
        self.mid_channels = self.in_channels * expand_ratio
        self.with_expand_conv = (self.mid_channels != self.in_channels)

        if self.with_se:
            assert isinstance(se_cfg, dict)

        if self.with_expand_conv:
            self.expand_conv = ConvModule(
                in_channels=self.in_channels,
                out_channels=self.mid_channels,
                kernel_size=1,
                stride=1,
                padding=0,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg)
        self.depthwise_conv = ConvModule(
            in_channels=self.mid_channels,
            out_channels=self.mid_channels,
            kernel_size=kernel_size,
            stride=self.stride,
            padding=kernel_size // 2,
            groups=self.mid_channels,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)
        if self.with_se:
            self.se = SELayer(self.mid_channels, **se_cfg)
        self.linear_conv = ConvModule(
            in_channels=self.mid_channels,
            out_channels=self.out_channels,
            kernel_size=1,
            stride=1,
            padding=0,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=None)

    def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor): The input tensor.
        Returns:
            torch.Tensor: The output tensor.
        """

        def _inner_forward(x):
            out = x

            if self.with_expand_conv:
                out = self.expand_conv(out)

            out = self.depthwise_conv(out)

            if self.with_se:
                out = self.se(out)

            out = self.linear_conv(out)

            if self.with_res_shortcut:
                return x + self.drop_path(out)
            else:
                return out

        if self.with_cp and x.requires_grad:
            out = cp.checkpoint(_inner_forward, x)
        else:
            out = _inner_forward(x)

        return out

mmrazor/models/ops/mobilenet_series.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
import torch.utils.checkpoint as cp
from mmcls.models.utils import channel_shuffle
from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule

from ..builder import OPS
from .base import BaseOP


@OPS.register_module()
class ShuffleBlock(BaseOP):
    """InvertedResidual block for Searchable ShuffleNetV2 backbone.

    Args:
        kernel_size (int): Size of the convolving kernel.
        stride (int): Stride of the convolution layer. Default: 1
        conv_cfg (dict, optional): Config dict for convolution layer.
            Default: None, which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='BN').
        act_cfg (dict): Config dict for activation layer.
            Default: dict(type='ReLU').
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.

    Returns:
        Tensor: The output tensor.
    """

    def __init__(self,
                 kernel_size,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 with_cp=False,
                 **kwargs):

        super(ShuffleBlock, self).__init__(**kwargs)

        assert kernel_size in [3, 5, 7]
        self.kernel_size = kernel_size
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.with_cp = with_cp

        branch_features = self.out_channels // 2
        if self.stride == 1:
            assert self.in_channels == branch_features * 2, (
                f'in_channels ({self.in_channels}) should equal to '
                f'branch_features * 2 ({branch_features * 2}) '
                'when stride is 1')

        if self.in_channels != branch_features * 2:
            assert self.stride != 1, (
                f'stride ({self.stride}) should not equal 1 when '
                f'in_channels != branch_features * 2')

        if self.stride > 1:
            self.branch1 = nn.Sequential(
                ConvModule(
                    self.in_channels,
                    self.in_channels,
                    kernel_size=self.kernel_size,
                    stride=self.stride,
                    padding=self.kernel_size // 2,
                    groups=self.in_channels,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    act_cfg=None),
                ConvModule(
                    self.in_channels,
                    branch_features,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    act_cfg=self.act_cfg),
            )

        self.branch2 = nn.Sequential(
            ConvModule(
                self.in_channels if (self.stride > 1) else branch_features,
                branch_features,
                kernel_size=1,
                stride=1,
                padding=0,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg),
            ConvModule(
                branch_features,
                branch_features,
                kernel_size=self.kernel_size,
                stride=self.stride,
                padding=self.kernel_size // 2,
                groups=branch_features,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                act_cfg=None),
            ConvModule(
                branch_features,
                branch_features,
                kernel_size=1,
                stride=1,
                padding=0,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg))

    def forward(self, x):

        def _inner_forward(x):
            if self.stride > 1:
                out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
            else:
                x1, x2 = x.chunk(2, dim=1)
                out = torch.cat((x1, self.branch2(x2)), dim=1)

            out = channel_shuffle(out, 2)

            return out

        if self.with_cp and x.requires_grad:
            out = cp.checkpoint(_inner_forward, x)
        else:
            out = _inner_forward(x)

        return out


@OPS.register_module()
class ShuffleXception(BaseOP):
    """Xception block for ShuffleNetV2 backbone.

    Args:
        conv_cfg (dict, optional): Config dict for convolution layer.
            Defaults to None, which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Defaults to dict(type='BN').
        act_cfg (dict): Config dict for activation layer.
            Defaults to dict(type='ReLU').
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Defaults to False.

    Returns:
        Tensor: The output tensor.
    """

    def __init__(self,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 with_cp=False,
                 **kwargs):
        super(ShuffleXception, self).__init__(**kwargs)
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.with_cp = with_cp
        self.mid_channels = self.out_channels // 2

        branch_features = self.out_channels // 2
        if self.stride == 1:
            assert self.in_channels == branch_features * 2, (
                f'in_channels ({self.in_channels}) should equal to '
                f'branch_features * 2 ({branch_features * 2}) '
                'when stride is 1')

        if self.in_channels != branch_features * 2:
            assert self.stride != 1, (
                f'stride ({self.stride}) should not equal 1 when '
                f'in_channels != branch_features * 2')

        if self.stride > 1:
            self.branch1 = nn.Sequential(
                ConvModule(
                    self.in_channels,
                    self.in_channels,
                    kernel_size=3,
                    stride=self.stride,
                    padding=1,
                    groups=self.in_channels,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    act_cfg=None),
                ConvModule(
                    self.in_channels,
                    branch_features,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    act_cfg=self.act_cfg),
            )

        self.branch2 = []

        self.branch2.append(
            DepthwiseSeparableConvModule(
                self.in_channels if (self.stride > 1) else branch_features,
                self.mid_channels,
                kernel_size=3,
                stride=self.stride,
                padding=1,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                dw_act_cfg=None,
                act_cfg=self.act_cfg), )
        self.branch2.append(
            DepthwiseSeparableConvModule(
                self.mid_channels,
                self.mid_channels,
                kernel_size=3,
                stride=1,
                padding=1,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                dw_act_cfg=None,
                act_cfg=self.act_cfg))
        self.branch2.append(
            DepthwiseSeparableConvModule(
                self.mid_channels,
                branch_features,
                kernel_size=3,
                stride=1,
                padding=1,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                dw_act_cfg=None,
                act_cfg=self.act_cfg))
        self.branch2 = nn.Sequential(*self.branch2)

    def forward(self, x):

        def _inner_forward(x):
            if self.stride > 1:
                out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
            else:
                x1, x2 = x.chunk(2, dim=1)
                out = torch.cat((x1, self.branch2(x2)), dim=1)

            out = channel_shuffle(out, 2)

            return out

        if self.with_cp and x.requires_grad:
            out = cp.checkpoint(_inner_forward, x)
        else:
            out = _inner_forward(x)

        return out

mmrazor/models/ops/shufflenet_series.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .ratio_pruning import RatioPruner
from .structure_pruning import StructurePruner
from .utils import *  # noqa: F401,F403

__all__ = ['RatioPruner', 'StructurePruner']

mmrazor/models/pruners/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
import torch.nn as nn
from torch.nn.modules import GroupNorm

from mmrazor.models.builder import PRUNERS
from .structure_pruning import StructurePruner
from .utils import SwitchableBatchNorm2d


@PRUNERS.register_module()
class RatioPruner(StructurePruner):
    """A random ratio pruner.

    Each layer can adjust its own width ratio randomly and independently.

    Args:
        ratios (list | tuple): Width ratio of each layer can be
            chosen from `ratios` randomly. The width ratio is the ratio between
            the number of reserved channels and that of all channels in a
            layer. For example, if `ratios` is [0.25, 0.5], there are 2 cases
            for us to choose from when we sample from a layer with 12 channels.
            One is sampling the very first 3 channels in this layer, another is
            sampling the very first 6 channels in this layer. Default to None.
    """

    def __init__(self, ratios, **kwargs):
        super(RatioPruner, self).__init__(**kwargs)
        ratios = list(ratios)
        ratios.sort()
        self.ratios = ratios
        self.min_ratio = ratios[0]

    def _check_pruner(self, supernet):
        for module in supernet.model.modules():
            if isinstance(module, GroupNorm):
                num_channels = module.num_channels
                num_groups = module.num_groups
                for ratio in self.ratios:
                    new_channels = int(round(num_channels * ratio))
                    assert (num_channels * ratio) % num_groups == 0, \
                        f'Expected number of channels in input of GroupNorm ' \
                        f'to be divisible by num_groups, but number of ' \
                        f'channels may be {new_channels} according to ' \
                        f'ratio {ratio} and num_groups={num_groups}'

    def prepare_from_supernet(self, supernet):
        super(RatioPruner, self).prepare_from_supernet(supernet)

    def get_channel_mask(self, out_mask):
        """Randomly choose a width ratio of a layer from ``ratios``"""
        out_channels = out_mask.size(1)
        random_ratio = np.random.choice(self.ratios)
        new_channels = int(round(out_channels * random_ratio))
        assert new_channels > 0, \
            'Output channels should be a positive integer.'
        new_out_mask = torch.zeros_like(out_mask)
        new_out_mask[:, :new_channels] = 1

        return new_out_mask

    def sample_subnet(self):
        """Random sample subnet by random mask.

        Returns:
            dict: Record the information to build the subnet from the supernet,
                its keys are the properties ``space_id`` in the pruner's search
                spaces, and its values are corresponding sampled out_mask.
        """
        subnet_dict = dict()
        for space_id, out_mask in self.channel_spaces.items():
            subnet_dict[space_id] = self.get_channel_mask(out_mask)
        return subnet_dict

    def set_min_channel(self):
        """Set the number of channels each layer to minimum."""
        subnet_dict = dict()
        for space_id, out_mask in self.channel_spaces.items():
            out_channels = out_mask.size(1)
            random_ratio = self.min_ratio
            new_channels = int(round(out_channels * random_ratio))
            assert new_channels > 0, \
                'Output channels should be a positive integer.'
            new_out_mask = torch.zeros_like(out_mask)
            new_out_mask[:, :new_channels] = 1

            subnet_dict[space_id] = new_out_mask

        self.set_subnet(subnet_dict)

    def switch_subnet(self, channel_cfg, subnet_ind=None):
        """Switch the channel config of the supernet according to channel_cfg.

        If we train more than one subnet together, we need to switch the
        channel_cfg from one to another during one training iteration.

        Args:
            channel_cfg (dict): The channel config of a subnet. Key is space_id
                and value is a dict which includes out_channels (and
                in_channels if exists).
            subnet_ind (int, optional): The index of the current subnet. If
                we replace normal BatchNorm2d with ``SwitchableBatchNorm2d``,
                we should switch the index of ``SwitchableBatchNorm2d`` when
                switch subnet. Defaults to None.
        """
        subnet_dict = dict()
        for name, channels_per_layer in channel_cfg.items():
            module = self.name2module[name]
            if (isinstance(module, SwitchableBatchNorm2d)
                    and subnet_ind is not None):
                # When switching bn we should switch index simultaneously
                module.index = subnet_ind
                continue

            out_channels = channels_per_layer['out_channels']
            out_mask = torch.zeros_like(module.out_mask)
            out_mask[:, :out_channels] = 1

            space_id = self.get_space_id(name)
            if space_id in subnet_dict:
                assert torch.equal(subnet_dict[space_id], out_mask)
            elif space_id is not None:
                subnet_dict[space_id] = out_mask

        self.set_subnet(subnet_dict)

    def convert_switchable_bn(self, module, num_bns):
        """Convert normal ``nn.BatchNorm2d`` to ``SwitchableBatchNorm2d``.

        Args:
            module (:obj:`torch.nn.Module`): The module to be converted.
            num_bns (int): The number of ``nn.BatchNorm2d`` in a
                ``SwitchableBatchNorm2d``.

        Return:
            :obj:`torch.nn.Module`: The converted module. Each
                ``nn.BatchNorm2d`` in this module has been converted to a
                ``SwitchableBatchNorm2d``.
        """
        module_output = module
        if isinstance(module, nn.modules.batchnorm._BatchNorm):
            module_output = SwitchableBatchNorm2d(module.num_features, num_bns)

        for name, child in module.named_children():
            module_output.add_module(
                name, self.convert_switchable_bn(child, num_bns))

        del module
        return module_output

mmrazor/models/pruners/ratio_pruning.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from abc import ABCMeta, abstractmethod
from collections import OrderedDict
from types import MethodType

import torch
import torch.nn as nn
from mmcv import digit_version
from mmcv.runner import BaseModule
from ordered_set import OrderedSet
from torch.nn.modules import GroupNorm
from torch.nn.modules.batchnorm import _BatchNorm
from torch.nn.modules.instancenorm import _InstanceNorm

from mmrazor.models.builder import PRUNERS
from .utils import SwitchableBatchNorm2d

# These grad_fn pattern are flags of specific a nn.Module
CONV = ('ThnnConv2DBackward', 'CudnnConvolutionBackward',
        'MkldnnConvolutionBackward')
FC = ('ThAddmmBackward', 'AddmmBackward', 'MmBackward')
BN = ('ThnnBatchNormBackward', 'CudnnBatchNormBackward',
      'NativeBatchNormBackward')
GN = ('NativeGroupNormBackward', )
CONCAT = ('CatBackward', )
# the modules which contains NON_PASS grad_fn need to change the parameter size
# according to channels after pruning
NON_PASS = CONV + FC
PASS = BN + GN
NORM = BN + GN

BACKWARD_PARSER_DICT = dict()
MAKE_GROUP_PARSER_DICT = dict()


def register_parser(parser_dict, name=None, force=False):
    """Register a parser function.

    A record will be added to ``parser_dict``, whose key is the specified
    ``name``, and value is the function itself.
    It can be used as a decorator or a normal function.

    Example:
        >>> BACKWARD_PARSER_DICT = dict()
        >>> @register_parser(BACKWARD_PARSER_DICT, 'ThnnConv2DBackward')
        >>> def conv_backward_parser():
        >>>     pass

    Args:
        parser_dict (dict): A dict to map strings to parser functions.
        name (str | None): The function name to be registered. If not
            specified, the function name will be used.
        force (bool, optional): Whether to override an existing function with
            the same name. Default: False.
    """

    def _register(parser_func):
        parser_name = parser_func.__name__ if name is None else name
        if (parser_name not in parser_dict) or force:
            parser_dict[parser_name] = parser_func
        else:
            raise KeyError(
                f'{parser_name} is already registered in task_dict, '
                'add "force=True" if you want to override it')
        return parser_func

    return _register


@PRUNERS.register_module()
class StructurePruner(BaseModule, metaclass=ABCMeta):
    """Base class for structure pruning. This class defines the basic functions
    of a structure pruner. Any pruner that inherits this class should at least
    define its own `sample_subnet` and `set_min_channel` functions. This part
    is being continuously optimized, and there may be major changes in the
    future.

    Reference to https://github.com/jshilong/FisherPruning

    Args:
        except_start_keys (List[str]): the module whose name start with a
            string in except_start_keys will not be prune.
    """

    def __init__(self, except_start_keys=['head.fc']):
        super(StructurePruner, self).__init__()
        if except_start_keys is None:
            self.except_start_keys = list()
        else:
            self.except_start_keys = except_start_keys

    def trace_shared_module_hook(self, module, inputs, outputs):
        """Trace shared modules. Modules such as the detection head in
        RetinaNet which are visited more than once during :func:`forward` are
        shared modules.

        Args:
            module (:obj:`torch.nn.Module`): The module to register hook.
            inputs (tuple): The input of the module.
            outputs (tuple): The output of the module.
        """
        module.cnt += 1
        if module.cnt == 2:
            self.shared_module.append(self.module2name[module])

    def prepare_from_supernet(self, supernet):
        """Prepare for pruning."""

        module2name = OrderedDict()
        name2module = OrderedDict()
        var2module = OrderedDict()

        # record the visited module name during trace path
        visited = dict()
        # Record shared modules which will be visited more than once during
        # forward such as shared detection head in RetinaNet.
        # If a module is not a shared module and it has been visited during
        # forward, its parent modules must have been traced already.
        # However, a shared module will be visited more than once during
        # forward, so it is still need to be traced even if it has been
        # visited.
        self.shared_module = []
        tmp_shared_module_hook_handles = list()

        for name, module in supernet.model.named_modules():
            if isinstance(module, nn.GroupNorm):
                min_required_version = '1.6.0'
                assert digit_version(
                    torch.__version__
                ) >= digit_version(min_required_version), (
                    f'Requires pytorch>={min_required_version} to auto-trace'
                    f'GroupNorm correctly.')
            if hasattr(module, 'weight'):
                # trace shared modules
                module.cnt = 0
                # the handle is only to remove the corresponding hook later
                handle = module.register_forward_hook(
                    self.trace_shared_module_hook)
                tmp_shared_module_hook_handles.append(handle)

                module2name[module] = name
                name2module[name] = module
                var2module[id(module.weight)] = module
                self.add_pruning_attrs(module)
                visited[name] = False
            if isinstance(module, SwitchableBatchNorm2d):
                name2module[name] = module
        self.name2module = name2module
        self.module2name = module2name

        # Set requires_grad to True. If the `requires_grad` of a module's
        # weight is False, we can not trace this module by parsing backward.
        param_require_grad = dict()
        for param in supernet.model.parameters():
            param_require_grad[id(param)] = param.requires_grad
            param.requires_grad = True

        pseudo_img = torch.randn(1, 3, 224, 224)
        # todo: support two stage detector and mmseg
        pseudo_img = supernet.forward_dummy(pseudo_img)
        pseudo_loss = supernet.cal_pseudo_loss(pseudo_img)

        # `trace_shared_module_hook` and `cnt` are only used to trace the
        # shared modules in a model and need to be remove later
        for name, module in supernet.model.named_modules():
            if hasattr(module, 'weight'):
                del module.cnt

        for handle in tmp_shared_module_hook_handles:
            handle.remove()

        # We set requires_grad to True to trace the whole architecture
        # topology. So it should be reset after that.
        for param in supernet.model.parameters():
            param.requires_grad = param_require_grad[id(param)]
        del param_require_grad

        non_pass_paths = list()
        cur_non_pass_path = list()
        self.trace_non_pass_path(pseudo_loss.grad_fn, module2name, var2module,
                                 cur_non_pass_path, non_pass_paths, visited)

        norm_conv_links = dict()
        self.trace_norm_conv_links(pseudo_loss.grad_fn, module2name,
                                   var2module, norm_conv_links, visited)
        self.norm_conv_links = norm_conv_links

        # a node can be the name of a conv module or a str like 'concat_{id}'
        node2parents = self.find_node_parents(non_pass_paths)
        self.node2parents = node2parents

        same_out_channel_groups = self.make_same_out_channel_groups(
            node2parents, name2module)

        self.module2group = dict()
        for group_name, group in same_out_channel_groups.items():
            for module_name in group:
                self.module2group[module_name] = group_name

        self.modules_have_ancest = list()
        for node_name, parents_name in node2parents.items():
            if node_name in name2module and len(parents_name) > 0:
                self.modules_have_ancest.append(node_name)

        self.modules_have_child = OrderedSet()
        for parents_name in node2parents.values():
            for name in parents_name:
                # The node is a module in supernet
                if name in name2module:
                    self.modules_have_child.add(name)

        self.channel_spaces = self.build_channel_spaces(name2module)

        self._reset_norm_running_stats(supernet)

    @abstractmethod
    def sample_subnet(self):
        """Sample a subnet from the supernet.

        Returns:
            dict: Record the information to build the subnet from the supernet,
                its keys are the properties ``space_id`` in the pruner's search
                spaces, and its values are corresponding sampled out_mask.
        """
        pass

    def get_space_id(self, module_name):
        """Get the corresponding space_id of the module_name.

        The modules who share the same space_id will share the same out_mask.
        If the module is the output module(there is no other ``nn.Module``
        whose input is its output), this function will return None. As the
        output module can not be pruned.
        If the input of this module is the concatenation of the output
        of several ``nn.Module``, this function will return a dict object.
        If this module is in one of the groups, this function will return the
        group name. As the modules in the same group should share the same
        space_id.
        Otherwise, this function will return the module_name as space_id.

        Args:
            module_name (str): the name of a ``nn.Module``.

        Return:
            str or dict or None: the corresponding space_id of the module_name.
        """
        if 'concat' in module_name and module_name not in self.name2module:
            # each module_name in concat_parents should be in name2module
            if 'item' in module_name:
                space_id = self.get_space_id(self.node2parents[module_name][0])
            else:
                concat_parents = [
                    self.get_space_id(parent)
                    for parent in self.node2parents[module_name]
                ]
                space_id = dict(concat=concat_parents)

        elif module_name not in self.modules_have_child:
            return None
        elif module_name in self.module2group:
            space_id = self.module2group[module_name]
        else:
            space_id = module_name
        return space_id

    def set_subnet(self, subnet_dict):
        """Modify the in_mask and out_mask of modules in supernet according to
        subnet_dict.

        Args:
            subnet_dict (dict): the key is space_id and the value is the
                corresponding sampled out_mask.
        """
        for module_name in self.modules_have_child:
            space_id = self.get_space_id(module_name)
            module = self.name2module[module_name]
            module.out_mask = subnet_dict[space_id].to(module.out_mask.device)

        for norm, conv in self.norm_conv_links.items():
            module = self.name2module[norm]
            conv_space_id = self.get_space_id(conv)
            # conv_space_id is None means the conv layer in front of
            # this normalization module can not be pruned. So we should not set
            # the out_mask of this normalization layer
            if conv_space_id is not None:
                module.out_mask = subnet_dict[conv_space_id].to(
                    module.out_mask.device)

        for module_name in self.modules_have_ancest:
            module = self.name2module[module_name]
            parents = self.node2parents[module_name]
            # To avoid ambiguity, we only allow the following two cases:
            # 1. all elements in parents are ``Conv2d``,
            # 2. there is only one element in parents, ``concat`` or ``chunk``
            # In case 1, all the ``Conv2d`` share the same space_id and
            # out_mask.
            # So in all cases, we only need the very first element in parents
            parent = parents[0]
            space_id = self.get_space_id(parent)

            if isinstance(space_id, dict):
                if 'concat' in space_id:
                    in_mask = []
                    for parent_space_id in space_id['concat']:
                        in_mask.append(subnet_dict[parent_space_id])
                    module.in_mask = torch.cat(
                        in_mask, dim=1).to(module.in_mask.device)
            else:
                module.in_mask = subnet_dict[space_id].to(
                    module.in_mask.device)

    def export_subnet(self):
        """Generate subnet configs according to the in_mask and out_mask of a
        module."""
        channel_cfg = dict()
        for name, module in self.name2module.items():

            channel_cfg[name] = dict()
            if hasattr(module, 'in_mask'):
                channel_cfg[name]['in_channels'] = int(
                    module.in_mask.cpu().sum())
                channel_cfg[name]['raw_in_channels'] = int(
                    module.in_mask.cpu().numel())

            if hasattr(module, 'out_mask'):
                channel_cfg[name]['out_channels'] = int(
                    module.out_mask.cpu().sum())
                channel_cfg[name]['raw_out_channels'] = int(
                    module.out_mask.cpu().numel())

        return channel_cfg

    def set_max_channel(self):
        """Set the number of channels each layer to maximum."""
        subnet_dict = dict()
        for space_id, out_mask in self.channel_spaces.items():
            new_out_mask = torch.ones_like(out_mask)
            subnet_dict[space_id] = new_out_mask
        self.set_subnet(subnet_dict)

    @abstractmethod
    def set_min_channel(self):
        """Set the number of channels each layer to minimum."""
        pass

    def find_make_group_parser(self, node_name, name2module):
        """Find the corresponding make_group_parser according to the
        ``node_name``"""
        if 'concat' in node_name and node_name not in name2module:
            return MAKE_GROUP_PARSER_DICT['concat']
        elif 'chunk' in node_name and node_name not in name2module:
            return MAKE_GROUP_PARSER_DICT['chunk']
        elif (node_name in name2module
              and isinstance(name2module[node_name], nn.Conv2d)
              and name2module[node_name].in_channels
              == name2module[node_name].out_channels
              and name2module[node_name].in_channels
              == name2module[node_name].groups):
            return MAKE_GROUP_PARSER_DICT['depthwise']
        else:
            return MAKE_GROUP_PARSER_DICT['common']

    @register_parser(MAKE_GROUP_PARSER_DICT, 'concat')
    def concat_make_group_parser(self, node_name, parents_name, group_idx,
                                 same_in_channel_groups,
                                 same_out_channel_groups):
        if 'item' in node_name:
            return self.make_group_parser(node_name, parents_name, group_idx,
                                          same_in_channel_groups,
                                          same_out_channel_groups)
        else:
            return group_idx, same_in_channel_groups, same_out_channel_groups

    @register_parser(MAKE_GROUP_PARSER_DICT, 'chunk')
    def chunk_make_group_parser(self, group_idx, same_in_channel_groups,
                                same_out_channel_groups, **kwargs):
        return group_idx, same_in_channel_groups, same_out_channel_groups

    @register_parser(MAKE_GROUP_PARSER_DICT, 'depthwise')
    def depthwise_make_group_parser(self, node_name, parents_name, **kwargs):
        # depth wise conv should be in the same group with its parent
        parents_name.add(node_name)
        return self.make_group_parser(node_name, parents_name, **kwargs)

    @register_parser(MAKE_GROUP_PARSER_DICT, 'common')
    def make_group_parser(self, node_name, parents_name, group_idx,
                          same_in_channel_groups, same_out_channel_groups):
        added = False
        for group_name in same_in_channel_groups:
            group_parents = set(same_out_channel_groups[group_name])
            if len(parents_name.intersection(group_parents)) > 0:
                same_in_channel_groups[group_name].append(node_name)
                same_out_channel_groups[group_name] = list(
                    parents_name.union(group_parents))
                added = True
                break
        if not added:
            group_idx += 1
            same_in_channel_groups[group_idx] = [node_name]
            same_out_channel_groups[group_idx] = list(parents_name)

        return group_idx, same_in_channel_groups, same_out_channel_groups

    def make_same_out_channel_groups(self, node2parents, name2module):
        """Modules have the same child should be in the same group."""
        idx = -1
        # the nodes in same_out_channel_groups are parents of
        # nodes in same_in_channel_groups
        same_in_channel_groups, same_out_channel_groups = {}, {}
        for node_name, parents_name in node2parents.items():
            parser = self.find_make_group_parser(node_name, name2module)
            idx, same_in_channel_groups, same_out_channel_groups = parser(
                self,
                node_name=node_name,
                parents_name=parents_name,
                group_idx=idx,
                same_in_channel_groups=same_in_channel_groups,
                same_out_channel_groups=same_out_channel_groups,
            )

        groups = dict()
        idx = 0
        for group in same_out_channel_groups.values():
            if len(group) > 1:
                group_name = f'group_{idx}'
                groups[group_name] = group
                idx += 1

        return groups

    @staticmethod
    def modify_conv_forward(module):
        """Modify the forward method of a conv layer."""
        original_forward = module.forward

        def modified_forward(self, feature):
            feature = feature * self.in_mask
            return original_forward(feature)

        return MethodType(modified_forward, module)

    @staticmethod
    def modify_fc_forward(module):
        """Modify the forward method of a linear layer."""
        original_forward = module.forward

        def modified_forward(self, feature):
            if not len(self.in_mask.shape) == len(self.out_mask.shape):
                self.in_mask = self.in_mask.reshape(self.in_mask.shape[:2])

            feature = feature * self.in_mask
            return original_forward(feature)

        return MethodType(modified_forward, module)

    def add_pruning_attrs(self, module):
        """Add masks to a ``nn.Module``."""
        if isinstance(module, nn.Conv2d):
            module.register_buffer(
                'in_mask',
                module.weight.new_ones((1, module.in_channels, 1, 1), ),
            )
            module.register_buffer(
                'out_mask',
                module.weight.new_ones((1, module.out_channels, 1, 1), ),
            )
            module.forward = self.modify_conv_forward(module)
        if isinstance(module, nn.Linear):
            module.register_buffer(
                'in_mask',
                module.weight.new_ones((1, module.in_features), ),
            )
            module.register_buffer(
                'out_mask',
                module.weight.new_ones((1, module.out_features), ),
            )
            module.forward = self.modify_fc_forward(module)
        if isinstance(module, _BatchNorm) or isinstance(
                module, _InstanceNorm) or isinstance(module, GroupNorm):
            module.register_buffer(
                'out_mask',
                module.weight.new_ones((1, len(module.weight), 1, 1), ),
            )

    def find_node_parents(self, paths):
        """Find the parent node of a node.

        A node in the ``paths`` can be a module name or a operation name such
        as `concat_140719322997152`. Note that the string of numbers following
        ``concat`` do not have a particular meaning. It just make the operation
        name unique.

        Args:
            paths (list): The traced paths.
        """
        node2parents = dict()
        for path in paths:
            if len(path) == 0:
                continue
            for i, node_name in enumerate(path[:-1]):
                parent_name = path[i + 1]
                if node_name in node2parents.keys():
                    node2parents[node_name].add(parent_name)
                else:
                    node2parents[node_name] = OrderedSet([parent_name])

            leaf_name = path[-1]
            if leaf_name not in node2parents.keys():
                node2parents[leaf_name] = OrderedSet()
        return node2parents

    def build_channel_spaces(self, name2module):
        """Build channel search space.

        Args:
            name2module (dict): A mapping between module_name and module.

        Return:
            dict: The channel search space. The key is space_id and the value
                is the corresponding out_mask.
        """
        search_space = dict()

        for module_name in self.modules_have_child:
            need_prune = True
            for key in self.except_start_keys:
                if module_name.startswith(key):
                    need_prune = False
                    break
            if not need_prune:
                continue
            if module_name in self.module2group:
                space_id = self.module2group[module_name]
            else:
                space_id = module_name
            module = name2module[module_name]
            if space_id not in search_space:
                search_space[space_id] = module.out_mask

        return search_space

    def deploy_subnet(self, supernet, channel_cfg):
        """Deploy subnet according `channel_cfg`."""
        for name, module in supernet.model.named_modules():
            if name not in channel_cfg:
                continue

            channels_per_layer = channel_cfg[name]
            requires_grad = module.weight.requires_grad
            out_channels = channels_per_layer['out_channels']
            temp_weight = module.weight.data[:out_channels]

            if hasattr(module, 'out_channels'):
                module.out_channels = out_channels
            if hasattr(module, 'out_features'):
                module.out_features = out_channels
            if hasattr(module, 'num_features'):
                module.num_features = out_channels
            if hasattr(module, 'out_mask'):
                module.out_mask = module.out_mask[:, :out_channels]

            if 'in_channels' in channels_per_layer:
                in_channels = channels_per_layer['in_channels']

                # can also handle depthwise conv
                temp_weight = temp_weight[:, :in_channels].data
                if hasattr(module, 'in_channels'):
                    module.in_channels = in_channels
                if hasattr(module, 'in_features'):
                    module.in_features = in_channels
                if hasattr(module, 'in_mask'):
                    module.in_mask = module.in_mask[:, :in_channels]
                # TODO Seems not support GroupConv
                if getattr(module, 'groups', in_channels) > 1:
                    module.groups = in_channels

            module.weight = nn.Parameter(temp_weight.data.clone())
            module.weight.requires_grad = requires_grad

            if hasattr(module, 'bias') and module.bias is not None:
                module.bias = nn.Parameter(
                    module.bias.data[:out_channels].clone())
                module.bias.requires_grad = requires_grad

            if hasattr(module, 'running_mean'):
                module.running_mean = module.running_mean[:out_channels]

            if hasattr(module, 'running_var'):
                module.running_var = module.running_var[:out_channels]

    def get_max_channel_bins(self, max_channel_bins):
        """Get the max number of channel bins of all the groups which can be
        pruned during searching.

        Args:
            max_channel_bins (int): The max number of bins in each layer.
        """
        channel_bins_dict = dict()
        for space_id in self.channel_spaces.keys():
            channel_bins_dict[space_id] = torch.ones((max_channel_bins, ))
        return channel_bins_dict

    def set_channel_bins(self, channel_bins_dict, max_channel_bins):
        """Set subnet according to the number of channel bins in a layer.

        Args:
            channel_bins_dict (dict): The number of bins in each layer. Key is
                the space_id of each layer and value is the corresponding
                mask of channel bin.
            max_channel_bins (int): The max number of bins in each layer.
        """
        subnet_dict = dict()
        for space_id, bin_mask in channel_bins_dict.items():
            mask = self.channel_spaces[space_id]
            shape = mask.shape
            channel_num = shape[1]
            channels_per_bin = channel_num // max_channel_bins
            new_mask = []
            for mask in bin_mask:
                new_mask.extend([1] * channels_per_bin if mask else [0] *
                                channels_per_bin)
            new_mask.extend([0] * (channel_num % max_channel_bins))
            new_mask = torch.tensor(new_mask).reshape(*shape)
            subnet_dict[space_id] = new_mask
        self.set_subnet(subnet_dict)

    def trace_non_pass_path(self, grad_fn, module2name, var2module, cur_path,
                            result_paths, visited):
        """Trace the topology of all the ``NON_PASS_MODULE``."""
        grad_fn = grad_fn[0] if isinstance(grad_fn, (list, tuple)) else grad_fn

        if grad_fn is not None:
            parser = self.find_backward_parser(grad_fn)
            if parser is not None:
                parser(self, grad_fn, module2name, var2module, cur_path,
                       result_paths, visited)
            else:
                # If the op is AccumulateGrad, parents is (),
                parents = grad_fn.next_functions
                if parents is not None:
                    for parent in parents:
                        self.trace_non_pass_path(parent, module2name,
                                                 var2module, cur_path,
                                                 result_paths, visited)
        else:
            result_paths.append(copy.deepcopy(cur_path))

    def trace_norm_conv_links(self, grad_fn, module2name, var2module,
                              norm_conv_links, visited):
        """Get the convolutional layer placed before a normalization layer in
        the model.

        Example:
            >>> conv = nn.Conv2d(3, 3, 3)
            >>> norm = nn.BatchNorm2d(3)
            >>> pseudo_img = torch.rand(1, 3, 224, 224)
            >>> out = norm(conv(pseudo_img))
            >>> print(out.grad_fn)
            <NativeBatchNormBackward object at 0x0000022BC709DB08>
            >>> print(out.grad_fn.next_functions)
            ((<ThnnConv2DBackward object at 0x0000020E40639688>, 0),
            (<AccumulateGrad object at 0x0000020E40639208>, 0),
            (<AccumulateGrad object at 0x0000020E406398C8>, 0))
            >>> # op.next_functions[0][0] is ThnnConv2DBackward means
            >>> # the parent of this NativeBatchNormBackward op is
            >>> # ThnnConv2DBackward
            >>> # op.next_functions[1][0].variable is the weight of this
            >>> # normalization module
            >>> # op.next_functions[2][0].variable is the bias of this
            >>> # normalization module

            >>> # Things are different in InstanceNorm
            >>> conv = nn.Conv2d(3, 3, 3)
            >>> norm = nn.InstanceNorm2d(3, affine=True)
            >>> out = norm(conv(pseudo_img))
            >>> print(out.grad_fn)
            <ViewBackward object at 0x0000022BC709DD48>
            >>> print(out.grad_fn.next_functions)
            ((<NativeBatchNormBackward object at 0x0000022BC81E8A08>, 0),)
            >>> print(out.grad_fn.next_functions[0][0].next_functions)
            ((<ViewBackward object at 0x0000022BC81E8DC8>, 0),
            (<RepeatBackward object at 0x0000022BC81E8D08>, 0),
            (<RepeatBackward object at 0x0000022BC81E81C8>, 0))
            >>> # Hence, a dfs is necessary.
        """

        def is_norm_grad_fn(grad_fn):
            for fn_name in NORM:
                if type(grad_fn).__name__.startswith(fn_name):
                    return True
            return False

        def is_conv_grad_fn(grad_fn):
            for fn_name in CONV:
                if type(grad_fn).__name__.startswith(fn_name):
                    return True
            return False

        def is_leaf_grad_fn(grad_fn):
            if type(grad_fn).__name__ == 'AccumulateGrad':
                return True
            return False

        grad_fn = grad_fn[0] if isinstance(grad_fn, (list, tuple)) else grad_fn
        if grad_fn is not None:
            if is_norm_grad_fn(grad_fn):
                conv_grad_fn = grad_fn.next_functions[0][0]
                while not is_conv_grad_fn(conv_grad_fn):
                    conv_grad_fn = conv_grad_fn.next_functions[0][0]

                leaf_grad_fn = conv_grad_fn.next_functions[1][0]
                while not is_leaf_grad_fn(leaf_grad_fn):
                    leaf_grad_fn = leaf_grad_fn.next_functions[0][0]
                conv_var = leaf_grad_fn.variable

                leaf_grad_fn = grad_fn.next_functions[1][0]
                while not is_leaf_grad_fn(leaf_grad_fn):
                    leaf_grad_fn = leaf_grad_fn.next_functions[0][0]
                bn_var = leaf_grad_fn.variable

                conv_module = var2module[id(conv_var)]
                bn_module = var2module[id(bn_var)]
                conv_name = module2name[conv_module]
                bn_name = module2name[bn_module]
                if visited[bn_name]:
                    pass
                else:
                    visited[bn_name] = True
                    norm_conv_links[bn_name] = conv_name

                    self.trace_norm_conv_links(conv_grad_fn, module2name,
                                               var2module, norm_conv_links,
                                               visited)

            else:
                # If the op is AccumulateGrad, parents is (),
                parents = grad_fn.next_functions
                if parents is not None:
                    for parent in parents:
                        self.trace_norm_conv_links(parent, module2name,
                                                   var2module, norm_conv_links,
                                                   visited)

    def find_backward_parser(self, grad_fn):
        for name, parser in BACKWARD_PARSER_DICT.items():
            if type(grad_fn).__name__.startswith(name):
                return parser

    @register_parser(BACKWARD_PARSER_DICT, 'ThnnConv2DBackward')
    @register_parser(BACKWARD_PARSER_DICT, 'CudnnConvolutionBackward')
    @register_parser(BACKWARD_PARSER_DICT, 'MkldnnConvolutionBackward')
    @register_parser(BACKWARD_PARSER_DICT, 'SlowConvDilated2DBackward')
    def conv_backward_parser(self, grad_fn, module2name, var2module, cur_path,
                             result_paths, visited):
        """Parse the backward of a conv layer.

        Example:
            >>> conv = nn.Conv2d(3, 3, 3)
            >>> pseudo_img = torch.rand(1, 3, 224, 224)
            >>> out = conv(pseudo_img)
            >>> print(out.grad_fn.next_functions)
            ((None, 0), (<AccumulateGrad object at 0x0000020E405CBD88>, 0),
            (<AccumulateGrad object at 0x0000020E405CB588>, 0))
            >>> # op.next_functions[0][0] is None means this ThnnConv2DBackward
            >>> # op has no parents
            >>> # op.next_functions[1][0].variable is the weight of this Conv2d
            >>> # module
            >>> # op.next_functions[2][0].variable is the bias of this Conv2d
            >>> # module
        """
        variable = grad_fn.next_functions[1][0].variable
        var_id = id(variable)
        module = var2module[var_id]
        name = module2name[module]
        parent = grad_fn.next_functions[0][0]
        cur_path.append(name)
        # If a module is not a shared module and it has been visited during
        # forward, its parent modules must have been traced already.
        # However, a shared module will be visited more than once during
        # forward, so it is still need to be traced even if it has been
        # visited.
        if visited[name] and name not in self.shared_module:
            result_paths.append(copy.deepcopy(cur_path))
        else:
            visited[name] = True
            self.trace_non_pass_path(parent, module2name, var2module, cur_path,
                                     result_paths, visited)
        cur_path.pop(-1)

    @register_parser(BACKWARD_PARSER_DICT, 'ThAddmmBackward')
    @register_parser(BACKWARD_PARSER_DICT, 'AddmmBackward')
    @register_parser(BACKWARD_PARSER_DICT, 'MmBackward')
    def linear_backward_parser(self, grad_fn, module2name, var2module,
                               cur_path, result_paths, visited):
        """Parse the backward of a conv layer.

        Example:
            >>> fc = nn.Linear(3, 3, bias=True)
            >>> input = torch.rand(3, 3)
            >>> out = fc(input)
            >>> print(out.grad_fn.next_functions)
            ((<AccumulateGrad object at 0x0000020E405F75C8>, 0), (None, 0),
            (<TBackward object at 0x0000020E405F7D48>, 0))
            >>> # op.next_functions[0][0].variable is the bias of this Linear
            >>> # module
            >>> # op.next_functions[1][0] is None means this AddmmBackward op
            >>> # has no parents
            >>> # op.next_functions[2][0] is the TBackward op, and
            >>> # op.next_functions[2][0].next_functions[0][0].variable is
            >>> # the transpose of the weight of this Linear module
        """
        variable = grad_fn.next_functions[2][0].next_functions[0][0].variable
        var_id = id(variable)
        module = var2module[var_id]
        name = module2name[module]
        parent = grad_fn.next_functions[1][0]
        cur_path.append(name)
        # If a module is not a shared module and it has been visited during
        # forward, its parent modules must have been traced already.
        # However, a shared module will be visited more than once during
        # forward, so it is still need to be traced even if it has been
        # visited.
        if visited[name] and name not in self.shared_module:
            result_paths.append(copy.deepcopy(cur_path))
        else:
            visited[name] = True
            self.trace_non_pass_path(parent, module2name, var2module, cur_path,
                                     result_paths, visited)
        cur_path.pop(-1)

    @register_parser(BACKWARD_PARSER_DICT, 'CatBackward')
    def concat_backward_parser(self, grad_fn, module2name, var2module,
                               cur_path, result_paths, visited):
        """Parse the backward of a concat operation.

        Example:
            >>> conv = nn.Conv2d(3, 3, 3)
            >>> pseudo_img = torch.rand(1, 3, 224, 224)
            >>> out1 = conv(pseudo_img)
            >>> out2 = conv(pseudo_img)
            >>> out = torch.cat([out1, out2], dim=1)
            >>> print(out.grad_fn.next_functions)
            ((<ThnnConv2DBackward object at 0x0000020E405F24C8>, 0),
            (<ThnnConv2DBackward object at 0x0000020E405F2648>, 0))
            >>> # the length of ``out.grad_fn.next_functions`` is two means
            >>> # ``out`` is obtained by concatenating two tensors
        """
        parents = grad_fn.next_functions
        concat_id = '_'.join([str(id(p)) for p in parents])
        name = f'concat_{concat_id}'
        cur_path.append(name)
        # If a module is not a shared module and it has been visited during
        # forward, its parent modules must have been traced already.
        # However, a shared module will be visited more than once during
        # forward, so it is still need to be traced even if it has been
        # visited.
        if (name in visited and visited[name]
                and name not in self.shared_module):
            result_paths.append(copy.deepcopy(cur_path))
        else:
            visited[name] = True
            for i, parent in enumerate(parents):
                cur_path.append(f'{name}_item_{i}')
                self.trace_non_pass_path(parent, module2name, var2module,
                                         cur_path, result_paths, visited)
                if cur_path.pop(-1) != f'{name}_item_{i}':
                    print(f'{name}_item_{i}')
        cur_path.pop(-1)

    @staticmethod
    def _reset_norm_running_stats(supernet):
        from torch.nn.modules.batchnorm import _NormBase

        for module in supernet.modules():
            if isinstance(module, _NormBase):
                module.reset_parameters()

mmrazor/models/pruners/structure_pruning.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .switchable_bn import SwitchableBatchNorm2d

__all__ = ['SwitchableBatchNorm2d']

mmrazor/models/pruners/utils/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn


class SwitchableBatchNorm2d(nn.Module):
    """Employs independent batch normalization for different switches in a
    slimmable network.

    To train slimmable networks, ``SwitchableBatchNorm2d`` privatizes all
    batch normalization layers for each switch in a slimmable network.
    Compared with the naive training approach, it solves the problem of feature
    aggregation inconsistency between different switches by independently
    normalizing the feature mean and variance during testing.

    Args:
        max_num_features (int): The maximum ``num_features`` among BatchNorm2d
            in all the switches.
        num_bns (int): The number of different switches in the slimmable
            networks.
    """

    def __init__(self, max_num_features, num_bns):
        super(SwitchableBatchNorm2d, self).__init__()

        self.max_num_features = max_num_features
        # number of BatchNorm2d in a SwitchableBatchNorm2d
        self.num_bns = num_bns
        bns = []
        for _ in range(num_bns):
            bns.append(nn.BatchNorm2d(max_num_features))
        self.bns = nn.ModuleList(bns)
        # When switching bn we should switch index simultaneously
        self.index = 0

    def forward(self, input):
        """Forward computation according to the current switch of the slimmable
        networks."""
        return self.bns[self.index](input)

mmrazor/models/pruners/utils/switchable_bn.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from .misc import find_latest_checkpoint
from .setup_env import setup_multi_processes

__all__ = ['find_latest_checkpoint', 'setup_multi_processes']

mmrazor/utils/__init__.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
norm_cfg = dict(type='BN',requires_grad=True)


student = dict(
    type = 'mmseg.EncoderDecoder',
    backbone=dict(
        type='mmseg.ResNetV1c',
        depth=18,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        dilations=(1, 1, 2, 4),
        strides=(1, 2, 1, 1),
        norm_cfg=norm_cfg,
        norm_eval=False,
        style='pytorch',
        contract_dilation=True),
    decode_head=dict(
        type='mmseg.PSPHead',
        in_channels=512,
        in_index=3,
        channels=128,
        pool_scales=(1, 2, 3, 6),
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='mmseg.CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)
    ),
    auxiliary_head=dict(
        type='mmseg.FCNHead',
        in_channels=256,
        in_index=2,
        channels=64,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='mmseg.CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)
    ),
    train_cfg=dict(),
    test_cfg=dict(mode='whole')
)


teacher = dict(
    type = 'mmseg.EncoderDecoder',
    # init_cfg=dict(
    #     type='Pretrained',
    #     checkpoint='pspnet_r101-d8_512x1024_80k_cityscapes_20200606_112211-e1e1100f.pth'),
    backbone=dict(
        type='mmseg.ResNetV1c',
        depth=101,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        dilations=(1, 1, 2, 4),
        strides=(1, 2, 1, 1),
        norm_cfg=norm_cfg,
        norm_eval=False,
        style='pytorch',
        contract_dilation=True
    ),
    
    decode_head=dict(
        type='mmseg.PSPHead',
        in_channels=2048,
        in_index=3,
        channels=512,
        pool_scales=(1, 2, 3, 6),
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='mmseg.CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)
    ),
)

algorithm = dict(
    type = 'GeneralDistill',
    architecture = dict(
        type = 'MMSegArchitecture',
        model = student,
    ),
    distiller = dict(
        type = 'SingleTeacherDistiller',
        teacher = teacher,
        teacher_trainable = False,
        components = [
            dict(
                student_module = 'decode_head.conv_seg',
                teacher_module = 'decode_head.conv_seg',
                losses = [
                    dict(
                        type='ChannelWiseDivergence',
                        name='loss_cwd_logits',
                        tau = 5,
                        loss_weight=3,
                    )
                ]
            )
        ]
    ),
    
)
tests/data/cwd_pspnet.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
norm_cfg = dict(type='SyncBN', requires_grad=True)
model = dict(
    type='mmdet.FasterRCNN',
    backbone=dict(
        type='mmcls.SearchableShuffleNetV2',
        norm_cfg=norm_cfg,
        out_indices=(0, 1, 2, 3),
        widen_factor=1.0,
        with_last_layer=False),
    neck=dict(
        type='FPN',
        norm_cfg=norm_cfg,
        in_channels=[64, 160, 320, 640],
        out_channels=256,
        num_outs=5),
    rpn_head=dict(
        type='RPNHead',
        in_channels=256,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            scales=[8],
            ratios=[0.5, 1.0, 2.0],
            strides=[4, 8, 16, 32, 64]),
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[.0, .0, .0, .0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        loss_cls=dict(
            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
    roi_head=dict(
        type='StandardRoIHead',
        bbox_roi_extractor=dict(
            type='SingleRoIExtractor',
            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
            out_channels=256,
            featmap_strides=[4, 8, 16, 32]),
        bbox_head=dict(
            type='Shared4Conv1FCBBoxHead',
            norm_cfg=norm_cfg,
            in_channels=256,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=80,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=False,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
            loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
    train_cfg=dict(
        rpn=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.7,
                neg_iou_thr=0.3,
                min_pos_iou=0.3,
                match_low_quality=True,
                ignore_iof_thr=-1),
            sampler=dict(
                type='RandomSampler',
                num=256,
                pos_fraction=0.5,
                neg_pos_ub=-1,
                add_gt_as_proposals=False),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        rpn_proposal=dict(
            nms_pre=2000,
            max_per_img=1000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.5,
                neg_iou_thr=0.5,
                min_pos_iou=0.5,
                match_low_quality=False,
                ignore_iof_thr=-1),
            sampler=dict(
                type='RandomSampler',
                num=512,
                pos_fraction=0.25,
                neg_pos_ub=-1,
                add_gt_as_proposals=True),
            pos_weight=-1,
            debug=False)),
    test_cfg=dict(
        rpn=dict(
            nms_pre=1000,
            max_per_img=1000,
            nms=dict(type='nms', iou_threshold=0.7),
            min_bbox_size=0),
        rcnn=dict(
            score_thr=0.05,
            nms=dict(type='nms', iou_threshold=0.5),
            max_per_img=100)
        # soft-nms is also supported for rcnn testing
        # e.g., nms=dict(type='soft_nms', iou_threshold=0.5, min_score=0.05)
    ),
)

mutator = dict(
    type='OneShotMutator',
    placeholder_mapping=dict(
        all_blocks=dict(
            type='OneShotOP',
            choices=dict(
                shuffle_3x3=dict(
                    type='ShuffleBlock', norm_cfg=norm_cfg, kernel_size=3),
                shuffle_5x5=dict(
                    type='ShuffleBlock', norm_cfg=norm_cfg, kernel_size=5),
                shuffle_7x7=dict(
                    type='ShuffleBlock', norm_cfg=norm_cfg, kernel_size=7),
                shuffle_xception=dict(type='ShuffleXception'),
            ))))

algorithm = dict(
    type='DetNAS',
    architecture=dict(
        type='MMDetArchitecture',
        model=model,
    ),
    mutator=mutator,
    pruner=None,
    distiller=None,
    retraining=False,
)


tests/data/detnas_frcnn_shufflenet_fpn.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import pytest


def test_v0_5_deprecated_set_random_seed() -> None:
    warn_msg = 'Deprecated in v0.3 and will be removed in v0.5, ' \
        'please import `set_random_seed` directly from `mmrazor.apis`'
    from mmrazor.apis.mmcls import set_random_seed
    with pytest.deprecated_call(match=warn_msg):
        set_random_seed(123)

    from mmrazor.apis.mmdet import set_random_seed
    with pytest.deprecated_call(match=warn_msg):
        set_random_seed(123)

    from mmrazor.apis.mmseg import set_random_seed
    with pytest.deprecated_call(match=warn_msg):
        set_random_seed(123)

tests/deprecated_api/test_remove_0-5.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from pathlib import Path

import mmcv
from mmcls.apis import inference_model
from mmdet.apis import inference_detector
from mmseg.apis import inference_segmentor

from mmrazor.apis import init_mmcls_model, init_mmdet_model, init_mmseg_model


def _sync_bn2bn(config: mmcv.Config) -> None:

    def dfs(cfg_dict) -> None:
        if isinstance(cfg_dict, dict):
            for k, v in cfg_dict.items():
                if k == 'norm_cfg':
                    if v['type'] == 'SyncBN':
                        v['type'] = 'BN'
                dfs(v)

    dfs(config._cfg_dict)


def test_init_mmcls_model() -> None:
    from mmcls.datasets import ImageNet

    config_file = 'configs/nas/spos/spos_subnet_shufflenetv2_8xb128_in1k.py'
    config = mmcv.Config.fromfile(config_file)
    config.model = None
    # Replace SyncBN with BN to inference on CPU
    _sync_bn2bn(config)

    mutable_file = 'configs/nas/spos/SPOS_SHUFFLENETV2_330M_IN1k_PAPER.yaml'
    model = init_mmcls_model(
        config,
        device='cpu',
        cfg_options={'algorithm.mutable_cfg': mutable_file})
    model.CLASSES = ImageNet.CLASSES
    assert not hasattr(model, 'architecture')
    assert hasattr(model, 'backbone')
    assert hasattr(model, 'neck')
    assert hasattr(model, 'head')

    img = mmcv.imread(Path(__file__).parent.parent / 'data/color.jpg', 'color')
    result = inference_model(model, img)
    assert isinstance(result, dict)
    assert result.get('pred_label') is not None
    assert result.get('pred_score') is not None
    assert result.get('pred_class') is not None


def test_init_mmdet_model() -> None:
    config_file = \
        'configs/nas/detnas/detnas_subnet_frcnn_shufflenetv2_fpn_1x_coco.py'
    config = mmcv.Config.fromfile(config_file)
    config.model = None
    # Replace SyncBN with BN to inference on CPU
    _sync_bn2bn(config)

    mutable_file = \
        'configs/nas/detnas/DETNAS_FRCNN_SHUFFLENETV2_340M_COCO_MMRAZOR.yaml'
    model = init_mmdet_model(
        config,
        device='cpu',
        cfg_options={'algorithm.mutable_cfg': mutable_file})
    assert not hasattr(model, 'architecture')

    img = mmcv.imread(Path(__file__).parent.parent / 'data/color.jpg', 'color')
    result = inference_detector(model, img)
    assert isinstance(result, list)


def test_init_mmseg_model() -> None:
    config_file = 'configs/distill/cwd/' \
        'cwd_cls_head_pspnet_r101_d8_pspnet_r18_d8_512x1024_cityscapes_80k.py'
    config = mmcv.Config.fromfile(config_file)
    config.model = None
    # Replace SyncBN with BN to inference on CPU
    _sync_bn2bn(config)

    # Enable test time augmentation
    config.data.test.pipeline[1].flip = True

    model = init_mmseg_model(config, device='cpu')
    assert not hasattr(model, 'architecture')
    assert hasattr(model, 'backbone')
    assert hasattr(model, 'decode_head')
    assert hasattr(model, 'auxiliary_head')

    img = mmcv.imread(Path(__file__).parent.parent / 'data/color.jpg', 'color')
    result = inference_segmentor(model, img)
    assert result[0].shape == (300, 400)

tests/test_apis/test_inference.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch

from mmrazor.apis import set_random_seed


def test_set_random_seed() -> None:
    set_random_seed(123, False)
    x1 = torch.rand(3, 3)
    x2 = np.random.rand(3, 3)

    set_random_seed(123, True)
    assert torch.backends.cudnn.deterministic
    assert not torch.backends.cudnn.benchmark
    y1 = torch.rand(3, 3)
    y2 = np.random.rand(3, 3)

    assert torch.allclose(x1, y1, 1e-6)
    assert np.allclose(x2, y2, 1e-6)

tests/test_apis/test_utils.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import MagicMock, patch

import torch
from torch.utils.data import DataLoader, Dataset

from mmrazor.core.builder import SEARCHERS


def collate_fn(data_batch):
    return data_batch


class ToyDataset(Dataset):
    METAINFO = dict()  # type: ignore
    data = torch.randn(12, 2)
    label = torch.ones(12)

    @property
    def metainfo(self):
        return self.METAINFO

    def __len__(self):
        return self.data.size(0)

    def __getitem__(self, index):
        return dict(inputs=self.data[index], data_sample=self.label[index])


class TestEvolutionSearcher(TestCase):

    def setUp(self):
        self.work_dir = tempfile.mkdtemp()
        self.searcher_cfg = dict(
            type='EvolutionSearcher',
            metrics='bbox',
            score_key='bbox_mAP',
            constraints=dict(flops=300 * 1e6),
            candidate_pool_size=50,
            candidate_top_k=10,
            max_epoch=20,
            num_mutation=20,
            num_crossover=20,
        )
        self.dataloader = DataLoader(ToyDataset(), collate_fn=collate_fn)
        self.test_fn = MagicMock()
        self.logger = MagicMock()

    def tearDown(self):
        shutil.rmtree(self.work_dir)

    @patch('mmrazor.models.algorithms.DetNAS')
    def test_init(self, mock_algo):
        cfg = copy.deepcopy(self.searcher_cfg)
        cfg['algorithm'] = mock_algo
        cfg['dataloader'] = self.dataloader
        cfg['test_fn'] = self.test_fn
        cfg['work_dir'] = self.work_dir
        cfg['logger'] = self.logger
        searcher = SEARCHERS.build(cfg)
        assert hasattr(searcher, 'algorithm')
        assert hasattr(searcher, 'logger')

        cfg['num_mutation'] = 40
        with self.assertRaises(ValueError):
            searcher = SEARCHERS.build(cfg)

tests/test_core/test_searcher.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import os
from copy import deepcopy
from os.path import dirname

import mmcv
import numpy as np
import torch
from mmcv import Config, ConfigDict

from mmrazor.models.builder import ALGORITHMS


def _demo_mm_inputs(input_shape=(1, 3, 8, 16), num_classes=10):
    """Create a superset of inputs needed to run test or train batches.

    Args:
        input_shape (tuple):
            input batch dimensions
        num_classes (int):
            number of semantic classes
    """
    (N, C, H, W) = input_shape

    rng = np.random.RandomState(0)

    imgs = rng.rand(*input_shape)
    segs = rng.randint(
        low=0, high=num_classes - 1, size=(N, 1, H, W)).astype(np.uint8)

    img_metas = [{
        'img_shape': (H, W, C),
        'ori_shape': (H, W, C),
        'pad_shape': (H, W, C),
        'filename': '<demo>.png',
        'scale_factor': 1.0,
        'flip': False,
        'flip_direction': 'horizontal'
    } for _ in range(N)]

    mm_inputs = {
        'img': torch.FloatTensor(imgs),
        'img_metas': img_metas,
        'gt_semantic_seg': torch.LongTensor(segs)
    }
    return mm_inputs


def test_autoslim_pretrain():
    model_cfg = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(
            type='mmcls.ResNet',
            depth=18,
            num_stages=4,
            out_indices=(3, ),
            style='pytorch'),
        neck=dict(type='mmcls.GlobalAveragePooling'),
        head=dict(
            type='mmcls.LinearClsHead',
            num_classes=1000,
            in_channels=512,
            loss=dict(type='mmcls.CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ))

    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=model_cfg,
    )

    pruner_cfg = dict(
        type='RatioPruner',
        ratios=[1 / 8, 2 / 8, 3 / 8, 4 / 8, 5 / 8, 6 / 8, 7 / 8, 1.0])

    distiller_cfg = dict(
        type='SelfDistiller',
        components=[
            dict(
                student_module='head.fc',
                teacher_module='head.fc',
                losses=[
                    dict(
                        type='KLDivergence',
                        name='loss_kd',
                        tau=1,
                        loss_weight=1,
                    )
                ]),
        ])

    algorithm_cfg = ConfigDict(
        type='AutoSlim',
        architecture=architecture_cfg,
        pruner=pruner_cfg,
        distiller=distiller_cfg)

    imgs = torch.randn(16, 3, 224, 224)
    label = torch.randint(0, 1000, (16, ))

    model = ALGORITHMS.build(algorithm_cfg)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    outputs = model.train_step({'img': imgs, 'gt_label': label}, optimizer)

    assert outputs['loss'].item() > 0
    assert outputs['num_samples'] == 16

    # test forward
    losses = model(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0


def test_autoslim_retrain():
    model_cfg = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(type='MobileNetV2', widen_factor=1.5),
        neck=dict(type='GlobalAveragePooling'),
        head=dict(
            type='LinearClsHead',
            num_classes=1000,
            in_channels=1920,
            loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ))

    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=model_cfg,
    )

    pruner_cfg = dict(
        type='RatioPruner',
        ratios=(2 / 12, 3 / 12, 4 / 12, 5 / 12, 6 / 12, 7 / 12, 8 / 12, 9 / 12,
                10 / 12, 11 / 12, 1.0))

    root_path = dirname(dirname(dirname(__file__)))
    channel_cfg = [
        os.path.join(root_path, 'data/AUTOSLIM_MBV2_530M_OFFICIAL.yaml'),
        os.path.join(root_path, 'data/AUTOSLIM_MBV2_320M_OFFICIAL.yaml')
    ]

    algorithm_cfg = ConfigDict(
        type='AutoSlim',
        architecture=architecture_cfg,
        pruner=pruner_cfg,
        retraining=True,
        channel_cfg=channel_cfg)

    imgs = torch.randn(16, 3, 224, 224)
    label = torch.randint(0, 1000, (16, ))

    # test multi subnet retraining
    model = ALGORITHMS.build(algorithm_cfg)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    outputs = model.train_step({'img': imgs, 'gt_label': label}, optimizer)

    assert outputs['loss'].item() > 0
    assert outputs['num_samples'] == 16

    # test single subnet retraining
    algorithm_cfg.channel_cfg = algorithm_cfg.channel_cfg[0]
    model = ALGORITHMS.build(algorithm_cfg)
    assert model.deployed
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    outputs = model.train_step({'img': imgs, 'gt_label': label}, optimizer)

    assert outputs['loss'].item() > 0
    assert outputs['num_samples'] == 16


def test_spos():

    model_cfg = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(type='SearchableShuffleNetV2', widen_factor=1.0),
        neck=dict(type='mmcls.GlobalAveragePooling'),
        head=dict(
            type='mmcls.LinearClsHead',
            num_classes=1000,
            in_channels=1024,
            loss=dict(type='mmcls.CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ),
    )

    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=model_cfg,
    )

    mutator_cfg = dict(
        type='OneShotMutator',
        placeholder_mapping=dict(
            all_blocks=dict(
                type='OneShotOP',
                choices=dict(
                    shuffle_3x3=dict(type='ShuffleBlock', kernel_size=3),
                    shuffle_5x5=dict(type='ShuffleBlock', kernel_size=5),
                    shuffle_7x7=dict(type='ShuffleBlock', kernel_size=7),
                    shuffle_xception=dict(type='ShuffleXception'),
                ))))

    algorithm_cfg = dict(
        type='SPOS',
        architecture=architecture_cfg,
        mutator=mutator_cfg,
        retraining=False,
    )

    imgs = torch.randn(16, 3, 224, 224)
    label = torch.randint(0, 1000, (16, ))
    spos_subnet_path = os.path.join(
        dirname(dirname(dirname(__file__))), 'data/spos_subnet.yaml')

    algorithm_cfg_ = deepcopy(algorithm_cfg)
    algorithm = ALGORITHMS.build(algorithm_cfg_)

    # test forward
    losses = algorithm(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0

    # test get_subnet_flops
    subnet_dict = algorithm.mutator.sample_subnet()
    flops_supernet = algorithm.get_subnet_flops()
    algorithm.mutator.set_subnet(subnet_dict)
    flops_subnet = algorithm.get_subnet_flops()
    assert flops_supernet > flops_subnet > 0

    # deploy_subnet in BaseMutator
    spos_subnet_dict = mmcv.fileio.load(spos_subnet_path)
    algorithm.mutator.deploy_subnet(algorithm.architecture, spos_subnet_dict)
    flops_subnet_spos = algorithm.get_subnet_flops()
    assert flops_supernet > flops_subnet_spos > 0


def test_spos_mobilenet():

    model_cfg = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(type='SearchableMobileNet', widen_factor=1.0),
        neck=dict(type='mmcls.GlobalAveragePooling'),
        head=dict(
            type='mmcls.LinearClsHead',
            num_classes=1000,
            in_channels=1280,
            loss=dict(type='mmcls.CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ),
    )

    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=model_cfg,
    )

    mutator_cfg = dict(
        type='OneShotMutator',
        placeholder_mapping=dict(
            searchable_blocks=dict(
                type='OneShotOP',
                choices=dict(
                    mb_k3e3=dict(
                        type='MBBlock',
                        kernel_size=3,
                        expand_ratio=3,
                        act_cfg=dict(type='ReLU6')),
                    mb_k5e3=dict(
                        type='MBBlock',
                        kernel_size=5,
                        expand_ratio=3,
                        act_cfg=dict(type='ReLU6')),
                    mb_k7e3=dict(
                        type='MBBlock',
                        kernel_size=7,
                        expand_ratio=3,
                        act_cfg=dict(type='ReLU6')),
                    mb_k3e6=dict(
                        type='MBBlock',
                        kernel_size=3,
                        expand_ratio=6,
                        act_cfg=dict(type='ReLU6')),
                    mb_k5e6=dict(
                        type='MBBlock',
                        kernel_size=5,
                        expand_ratio=6,
                        act_cfg=dict(type='ReLU6')),
                    mb_k7e6=dict(
                        type='MBBlock',
                        kernel_size=7,
                        expand_ratio=6,
                        act_cfg=dict(type='ReLU6')),
                    identity=dict(type='Identity'))),
            first_blocks=dict(
                type='OneShotOP',
                choices=dict(
                    mb_k3e1=dict(
                        type='MBBlock',
                        kernel_size=3,
                        expand_ratio=1,
                        act_cfg=dict(type='ReLU6')), ))))

    algorithm_cfg = dict(
        type='SPOS',
        architecture=architecture_cfg,
        mutator=mutator_cfg,
        retraining=False,
    )

    imgs = torch.randn(16, 3, 224, 224)
    label = torch.randint(0, 1000, (16, ))

    algorithm_cfg_ = deepcopy(algorithm_cfg)
    algorithm = ALGORITHMS.build(algorithm_cfg_)

    # test forward
    losses = algorithm(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0


def test_detnas():
    config_path = os.path.join(
        dirname(dirname(dirname(__file__))),
        'data/detnas_frcnn_shufflenet_fpn.py')
    config = Config.fromfile(config_path)

    # test detnas init
    algorithm = ALGORITHMS.build(config.algorithm)
    assert hasattr(algorithm, 'architecture')


def test_darts():

    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=dict(
            type='mmcls.ImageClassifier',
            backbone=dict(
                type='DartsBackbone',
                in_channels=3,
                base_channels=8,
                num_layers=4,
                num_nodes=4,
                stem_multiplier=3,
                out_indices=(3, )),
            neck=dict(type='GlobalAveragePooling'),
            head=dict(
                type='LinearClsHead',
                num_classes=10,
                in_channels=128,
                loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
                topk=(1, 5),
            ),
        ),
    )

    mutator_cfg = dict(
        type='DartsMutator',
        placeholder_mapping=dict(
            node=dict(
                type='DifferentiableOP',
                with_arch_param=True,
                choices=dict(
                    zero=dict(type='DartsZero'),
                    skip_connect=dict(
                        type='DartsSkipConnect',
                        norm_cfg=dict(type='BN', affine=False)),
                    max_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='max',
                        norm_cfg=dict(type='BN', affine=False)),
                    avg_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='avg',
                        norm_cfg=dict(type='BN', affine=False)),
                    sep_conv_3x3=dict(
                        type='DartsSepConv',
                        kernel_size=3,
                        norm_cfg=dict(type='BN', affine=False)),
                    sep_conv_5x5=dict(
                        type='DartsSepConv',
                        kernel_size=5,
                        norm_cfg=dict(type='BN', affine=False)),
                    dil_conv_3x3=dict(
                        type='DartsDilConv',
                        kernel_size=3,
                        norm_cfg=dict(type='BN', affine=False)),
                    dil_conv_5x5=dict(
                        type='DartsDilConv',
                        kernel_size=5,
                        norm_cfg=dict(type='BN', affine=False)),
                )),
            node_edge=dict(
                type='DifferentiableEdge',
                num_chosen=2,
                with_arch_param=False,
            )),
    )

    algorithm_cfg = dict(
        type='Darts',
        architecture=architecture_cfg,
        mutator=mutator_cfg,
        retraining=False,
        unroll=False,
    )

    imgs = torch.randn(16, 3, 32, 32)
    label = torch.randint(0, 10, (16, ))

    algorithm_cfg_ = deepcopy(algorithm_cfg)
    algorithm = ALGORITHMS.build(algorithm_cfg_)

    optimizers = dict(
        architecture=torch.optim.SGD(
            algorithm.architecture.parameters(), lr=0.01),
        mutator=torch.optim.SGD(algorithm.mutator.parameters(), lr=0.01),
    )

    data = [{'img': imgs, 'gt_label': label}, {'img': imgs, 'gt_label': label}]
    # test forward
    losses = algorithm.train_step(data, optimizers)
    assert losses['loss'].item() > 0


def test_cwd():
    config_path = './tests/data/cwd_pspnet.py'

    config = Config.fromfile(config_path)

    mm_inputs = _demo_mm_inputs(input_shape=(16, 3, 256, 256), num_classes=19)
    algorithm = ALGORITHMS.build(config.algorithm)

    # test algorithm train_step
    losses = algorithm.train_step(mm_inputs, None)
    assert losses['loss'].item() > 0


def test_rkd():
    student = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(
            type='ResNet',
            depth=18,
            num_stages=4,
            out_indices=(3, ),
            style='pytorch'),
        neck=dict(type='GlobalAveragePooling'),
        head=dict(
            type='LinearClsHead',
            num_classes=1000,
            in_channels=512,
            loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ))

    teacher = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(
            type='ResNet',
            depth=34,
            num_stages=4,
            out_indices=(3, ),
            style='pytorch'),
        neck=dict(type='GlobalAveragePooling'),
        head=dict(
            type='LinearClsHead',
            num_classes=1000,
            in_channels=512,
            loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ))

    # test RelationalKD w/ l2 norm
    algorithm_cfg = ConfigDict(
        type='GeneralDistill',
        architecture=dict(
            type='MMClsArchitecture',
            model=student,
        ),
        with_student_loss=True,
        with_teacher_loss=False,
        distiller=dict(
            type='SingleTeacherDistiller',
            teacher=teacher,
            teacher_trainable=False,
            teacher_norm_eval=True,
            components=[
                dict(
                    student_module='neck.gap',
                    teacher_module='neck.gap',
                    losses=[
                        dict(
                            type='DistanceWiseRKD',
                            name='distance_wise_loss',
                            loss_weight=25.0,
                            with_l2_norm=True),
                        dict(
                            type='AngleWiseRKD',
                            name='angle_wise_loss',
                            loss_weight=50.0,
                            with_l2_norm=True),
                    ])
            ]),
    )

    imgs = torch.randn(16, 3, 32, 32)
    label = torch.randint(0, 10, (16, ))

    algorithm = ALGORITHMS.build(algorithm_cfg)

    optimizer = torch.optim.SGD(algorithm.parameters(), lr=0.01)
    outputs = algorithm.train_step({'img': imgs, 'gt_label': label}, optimizer)
    assert outputs['loss'].item() > 0
    assert outputs['num_samples'] == 16

    # test forward
    losses = algorithm(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0

    # test RelationalKD w/o l2 norm
    algorithm_cfg.distiller.components = [
        dict(
            student_module='neck.gap',
            teacher_module='neck.gap',
            losses=[
                dict(
                    type='DistanceWiseRKD',
                    name='distance_wise_loss',
                    loss_weight=25.0,
                    with_l2_norm=False),
                dict(
                    type='AngleWiseRKD',
                    name='angle_wise_loss',
                    loss_weight=50.0,
                    with_l2_norm=False),
            ])
    ]

    optimizer = torch.optim.SGD(algorithm.parameters(), lr=0.01)
    outputs = algorithm.train_step({'img': imgs, 'gt_label': label}, optimizer)
    assert outputs['loss'].item() > 0
    assert outputs['num_samples'] == 16

    losses = algorithm(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0

tests/test_models/test_algorithms/test_algorithm.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from mmseg.models import BACKBONES, HEADS
from mmseg.models.decode_heads.cascade_decode_head import BaseCascadeDecodeHead
from mmseg.models.decode_heads.decode_head import BaseDecodeHead
from torch import nn


def _demo_mm_inputs(input_shape=(1, 3, 8, 16), num_classes=10):
    """Create a superset of inputs needed to run test or train batches.

    Args:
        input_shape (tuple):
            input batch dimensions
        num_classes (int):
            number of semantic classes
    """
    (N, C, H, W) = input_shape

    rng = np.random.RandomState(0)

    imgs = rng.rand(*input_shape)
    segs = rng.randint(
        low=0, high=num_classes - 1, size=(N, 1, H, W)).astype(np.uint8)

    img_metas = [{
        'img_shape': (H, W, C),
        'ori_shape': (H, W, C),
        'pad_shape': (H, W, C),
        'filename': '<demo>.png',
        'scale_factor': 1.0,
        'flip': False,
        'flip_direction': 'horizontal'
    } for _ in range(N)]

    mm_inputs = {
        'imgs': torch.FloatTensor(imgs),
        'img_metas': img_metas,
        'gt_semantic_seg': torch.LongTensor(segs)
    }
    return mm_inputs


@BACKBONES.register_module()
class ExampleBackbone(nn.Module):

    def __init__(self):
        super(ExampleBackbone, self).__init__()
        self.conv = nn.Conv2d(3, 3, 3)

    def init_weights(self, pretrained=None):
        pass

    def forward(self, x):
        return [self.conv(x)]


@HEADS.register_module()
class ExampleDecodeHead(BaseDecodeHead):

    def __init__(self):
        super(ExampleDecodeHead, self).__init__(3, 3, num_classes=19)

    def forward(self, inputs):
        return self.cls_seg(inputs[0])


@HEADS.register_module()
class ExampleCascadeDecodeHead(BaseCascadeDecodeHead):

    def __init__(self):
        super(ExampleCascadeDecodeHead, self).__init__(3, 3, num_classes=19)

    def forward(self, inputs, prev_out):
        return self.cls_seg(inputs[0])


def _segmentor_forward_train_test(segmentor):
    if isinstance(segmentor.decode_head, nn.ModuleList):
        num_classes = segmentor.decode_head[-1].num_classes
    else:
        num_classes = segmentor.decode_head.num_classes
    # batch_size=2 for BatchNorm
    mm_inputs = _demo_mm_inputs(num_classes=num_classes)

    imgs = mm_inputs.pop('imgs')
    img_metas = mm_inputs.pop('img_metas')
    gt_semantic_seg = mm_inputs['gt_semantic_seg']

    # convert to cuda Tensor if applicable
    if torch.cuda.is_available():
        segmentor = segmentor.cuda()
        imgs = imgs.cuda()
        gt_semantic_seg = gt_semantic_seg.cuda()

    # Test forward train
    losses = segmentor.forward(
        imgs, img_metas, gt_semantic_seg=gt_semantic_seg, return_loss=True)
    assert isinstance(losses, dict)

    # Test train_step
    data_batch = dict(
        img=imgs, img_metas=img_metas, gt_semantic_seg=gt_semantic_seg)
    outputs = segmentor.train_step(data_batch, None)
    assert isinstance(outputs, dict)
    assert 'loss' in outputs
    assert 'log_vars' in outputs
    assert 'num_samples' in outputs

    # Test val_step
    with torch.no_grad():
        segmentor.eval()
        data_batch = dict(
            img=imgs, img_metas=img_metas, gt_semantic_seg=gt_semantic_seg)
        outputs = segmentor.val_step(data_batch, None)
        assert isinstance(outputs, dict)
        assert 'loss' in outputs
        assert 'log_vars' in outputs
        assert 'num_samples' in outputs

    # Test forward simple test
    with torch.no_grad():
        segmentor.eval()
        # pack into lists
        img_list = [img[None, :] for img in imgs]
        img_meta_list = [[img_meta] for img_meta in img_metas]
        segmentor.forward(img_list, img_meta_list, return_loss=False)

    # Test forward aug test
    with torch.no_grad():
        segmentor.eval()
        # pack into lists
        img_list = [img[None, :] for img in imgs]
        img_list = img_list + img_list
        img_meta_list = [[img_meta] for img_meta in img_metas]
        img_meta_list = img_meta_list + img_meta_list
        segmentor.forward(img_list, img_meta_list, return_loss=False)

tests/test_models/test_algorithms/utils.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import tempfile
from copy import deepcopy

import numpy as np
import torch

from mmrazor.models.builder import ARCHITECTURES


def test_architecture_mmcls():
    model_cfg = dict(
        dict(
            type='mmcls.ImageClassifier',
            backbone=dict(
                type='mmcls.ResNet_CIFAR',
                depth=50,
                num_stages=4,
                out_indices=(3, ),
                style='pytorch'),
            neck=dict(type='mmcls.GlobalAveragePooling'),
            head=dict(
                type='mmcls.LinearClsHead',
                num_classes=10,
                in_channels=2048,
                loss=dict(type='CrossEntropyLoss'))), )

    architecture_cfg = dict(type='MMClsArchitecture', model=model_cfg)

    imgs = torch.randn(16, 3, 32, 32)
    label = torch.randint(0, 10, (16, ))

    supernet_cfg_ = deepcopy(architecture_cfg)
    architecture = ARCHITECTURES.build(supernet_cfg_)

    # test property
    assert architecture.model.with_neck
    assert architecture.model.with_head

    # test train_step
    outputs = architecture.model.train_step({
        'img': imgs,
        'gt_label': label
    }, None)
    assert outputs['loss'].item() > 0
    assert outputs['num_samples'] == 16

    # test val_step
    outputs = architecture.model.val_step({
        'img': imgs,
        'gt_label': label
    }, None)
    assert outputs['loss'].item() > 0
    assert outputs['num_samples'] == 16

    # test forward
    losses = architecture(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0

    # test forward_test
    architecture_cfg_ = deepcopy(architecture_cfg)
    architecture = ARCHITECTURES.build(architecture_cfg_)
    pred = architecture(imgs, return_loss=False, img_metas=None)
    assert isinstance(pred, list) and len(pred) == 16

    single_img = torch.randn(1, 3, 32, 32)
    pred = architecture(single_img, return_loss=False, img_metas=None)
    assert isinstance(pred, list) and len(pred) == 1

    # test simple_test
    single_img = torch.randn(1, 3, 32, 32)
    pred = architecture.simple_test(single_img, img_metas=None)
    assert isinstance(pred, list) and len(pred) == 1

    # test show_result
    img = np.random.randint(0, 255, (224, 224, 3)).astype(np.uint8)
    result = dict(pred_class='cat', pred_label=0, pred_score=0.9)

    with tempfile.TemporaryDirectory() as tmpdir:
        out_file = osp.join(tmpdir, 'out.png')
        architecture.show_result(img, result, out_file=out_file)
        assert osp.exists(out_file)

    with tempfile.TemporaryDirectory() as tmpdir:
        out_file = osp.join(tmpdir, 'out.png')
        architecture.show_result(img, result, out_file=out_file)
        assert osp.exists(out_file)

tests/test_models/test_architecture.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch

from mmrazor.models.builder import MUTABLES


def test_one_shot_op():
    oneshot_choice_op = dict(
        type='OneShotOP',
        space_id='test',
        num_chosen=1,
        choices=dict(
            shuffle_3x3=dict(type='ShuffleBlock', kernel_size=3),
            shuffle_5x5=dict(type='ShuffleBlock', kernel_size=5),
            shuffle_7x7=dict(type='ShuffleBlock', kernel_size=7),
            shuffle_xception=dict(type='ShuffleXception'),
        ),
        choice_args=dict(in_channels=16, out_channels=16, stride=1))

    model = MUTABLES.build(oneshot_choice_op)

    tensor = torch.randn(16, 16, 32, 32)

    # test forward
    outputs = model(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32


def test_differentiable_op():
    oneshot_choice_op = dict(
        type='DifferentiableOP',
        space_id='test',
        num_chosen=1,
        with_arch_param=True,
        choices=dict(
            zero=dict(type='DartsZero'),
            skip_connect=dict(type='DartsSkipConnect'),
            dil_conv_3x3=dict(type='DartsDilConv', kernel_size=3),
            dil_conv_5x5=dict(type='DartsDilConv', kernel_size=5),
            sep_conv_3x3=dict(type='DartsSepConv', kernel_size=3),
            sep_conv_5x5=dict(type='DartsSepConv', kernel_size=5),
        ),
        choice_args=dict(in_channels=16, out_channels=16, stride=2))

    model = MUTABLES.build(oneshot_choice_op)
    arch_param = model.build_arch_param()

    tensor = torch.randn(16, 16, 32, 32)

    # test forward
    outputs = model(tensor, arch_param)
    assert outputs.size(1) == 16 and outputs.size(2) == 16

tests/test_models/test_mutable.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from copy import deepcopy

import torch

from mmrazor.models.builder import ARCHITECTURES, MUTATORS


def test_one_shot_mutator():
    model_cfg = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(type='SearchableShuffleNetV2', widen_factor=1.0),
        neck=dict(type='mmcls.GlobalAveragePooling'),
        head=dict(
            type='mmcls.LinearClsHead',
            num_classes=1000,
            in_channels=1024,
            loss=dict(type='mmcls.CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ),
    )

    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=model_cfg,
    )

    mutator_cfg = dict(
        type='OneShotMutator',
        placeholder_mapping=dict(
            all_blocks=dict(
                type='OneShotOP',
                choices=dict(
                    shuffle_3x3=dict(type='ShuffleBlock', kernel_size=3),
                    shuffle_5x5=dict(type='ShuffleBlock', kernel_size=5),
                    shuffle_7x7=dict(type='ShuffleBlock', kernel_size=7),
                    shuffle_xception=dict(type='ShuffleXception'),
                ))))

    imgs = torch.randn(16, 3, 224, 224)
    label = torch.randint(0, 1000, (16, ))

    architecture_cfg_ = deepcopy(architecture_cfg)
    architecture = ARCHITECTURES.build(architecture_cfg_)

    mutator_cfg_ = deepcopy(mutator_cfg)
    mutator = MUTATORS.build(mutator_cfg_)

    mutator.prepare_from_supernet(architecture)
    assert hasattr(mutator, 'search_spaces')
    assert len(mutator.search_spaces) > 0

    # test forward
    subnet_dict = mutator.sample_subnet()
    mutator.set_subnet(subnet_dict)
    losses = architecture(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0

    # test set_chosen_subnet
    mutator.set_chosen_subnet(subnet_dict)
    search_spaces = mutator.search_spaces
    for key in search_spaces.keys():
        assert 'chosen' in search_spaces[key].keys()

    # test mutation and crossover
    subnet_dict1 = mutator.sample_subnet()
    subnet_dict2 = mutator.sample_subnet()
    mutation_subnet_dict = mutator.mutation(subnet_dict1)
    crossover_subnet_dict = mutator.crossover(subnet_dict1, subnet_dict2)
    assert isinstance(mutation_subnet_dict, dict)
    assert len(mutation_subnet_dict) > 0
    assert isinstance(crossover_subnet_dict, dict)
    assert len(crossover_subnet_dict) > 0


def test_differentiable_mutator():
    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=dict(
            type='mmcls.ImageClassifier',
            backbone=dict(
                type='DartsBackbone',
                in_channels=3,
                base_channels=8,
                num_layers=4,
                num_nodes=4,
                stem_multiplier=3,
                out_indices=(3, )),
            neck=dict(type='GlobalAveragePooling'),
            head=dict(
                type='LinearClsHead',
                num_classes=10,
                in_channels=128,
                loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
                topk=(1, 5),
            ),
        ),
    )

    mutator_cfg = dict(
        type='DartsMutator',
        placeholder_mapping=dict(
            node=dict(
                type='DifferentiableOP',
                with_arch_param=True,
                choices=dict(
                    zero=dict(type='DartsZero'),
                    skip_connect=dict(
                        type='DartsSkipConnect',
                        norm_cfg=dict(type='BN', affine=False)),
                    max_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='max',
                        norm_cfg=dict(type='BN', affine=False)),
                    avg_pool_3x3=dict(
                        type='DartsPoolBN',
                        pool_type='avg',
                        norm_cfg=dict(type='BN', affine=False)),
                    sep_conv_3x3=dict(
                        type='DartsSepConv',
                        kernel_size=3,
                        norm_cfg=dict(type='BN', affine=False)),
                    sep_conv_5x5=dict(
                        type='DartsSepConv',
                        kernel_size=5,
                        norm_cfg=dict(type='BN', affine=False)),
                    dil_conv_3x3=dict(
                        type='DartsDilConv',
                        kernel_size=3,
                        norm_cfg=dict(type='BN', affine=False)),
                    dil_conv_5x5=dict(
                        type='DartsDilConv',
                        kernel_size=5,
                        norm_cfg=dict(type='BN', affine=False)),
                )),
            node_edge=dict(
                type='DifferentiableEdge',
                num_chosen=2,
                with_arch_param=False,
            )),
    )

    imgs = torch.randn(16, 3, 32, 32)
    label = torch.randint(0, 10, (16, ))

    architecture_cfg_ = deepcopy(architecture_cfg)
    architecture = ARCHITECTURES.build(architecture_cfg_)

    mutator_cfg_ = deepcopy(mutator_cfg)
    mutator = MUTATORS.build(mutator_cfg_)

    # test prepare_from_superbet
    mutator.prepare_from_supernet(architecture)
    assert hasattr(mutator, 'search_spaces')
    assert hasattr(mutator, 'arch_params')

    # test forward
    losses = architecture(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0

    # test search_subnet
    # TODO

tests/test_models/test_mutator.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import torch

from mmrazor.models.builder import OPS


def test_common_ops():
    tensor = torch.randn(16, 16, 32, 32)

    # test stride != 1
    identity_cfg = dict(
        type='Identity', in_channels=16, out_channels=16, stride=2)

    op = OPS.build(identity_cfg)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 16

    # test stride == 1
    identity_cfg = dict(
        type='Identity', in_channels=16, out_channels=16, stride=1)

    op = OPS.build(identity_cfg)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test in_channels != out_channels
    identity_cfg = dict(
        type='Identity', in_channels=8, out_channels=16, stride=1)

    op = OPS.build(identity_cfg)

    # test forward
    outputs = op(tensor[:, :8])
    assert outputs.size(1) == 16 and outputs.size(2) == 32


def test_shuffle_series():

    tensor = torch.randn(16, 16, 32, 32)

    # test ShuffleBlock_7x7
    shuffle_block_7x7 = dict(
        type='ShuffleBlock',
        in_channels=16,
        out_channels=16,
        kernel_size=7,
        stride=1)

    op = OPS.build(shuffle_block_7x7)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test ShuffleBlock_5x5
    shuffle_block_5x5 = dict(
        type='ShuffleBlock',
        in_channels=16,
        out_channels=16,
        kernel_size=5,
        stride=1)

    op = OPS.build(shuffle_block_5x5)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test ShuffleBlock_3x3
    shuffle_block_3x3 = dict(
        type='ShuffleBlock',
        in_channels=16,
        out_channels=16,
        kernel_size=3,
        stride=1)

    op = OPS.build(shuffle_block_3x3)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test ShuffleXception
    shuffle_xception = dict(
        type='ShuffleXception', in_channels=16, out_channels=16, stride=1)

    op = OPS.build(shuffle_xception)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32


def test_mobilenet_series():

    tensor = torch.randn(16, 16, 32, 32)

    kernel_sizes = (3, 5, 7)
    expand_ratios = (3, 6)
    strides = (1, 2)
    se_cfg_1 = dict(
        ratio=4,
        act_cfg=(dict(type='HSwish'),
                 dict(
                     type='HSigmoid',
                     bias=3,
                     divisor=6,
                     min_value=0,
                     max_value=1)))
    se_cfgs = (None, se_cfg_1)
    drop_path_rates = (0, 0.2)
    with_cps = (True, False)

    for kernel_size in kernel_sizes:
        for expand_ratio in expand_ratios:
            for stride in strides:
                for se_cfg in se_cfgs:
                    for drop_path_rate in drop_path_rates:
                        for with_cp in with_cps:
                            op_cfg = dict(
                                type='MBBlock',
                                in_channels=16,
                                out_channels=16,
                                kernel_size=kernel_size,
                                expand_ratio=expand_ratio,
                                se_cfg=se_cfg,
                                drop_path_rate=drop_path_rate,
                                with_cp=with_cp,
                                stride=stride)

                    op = OPS.build(op_cfg)

                    # test forward
                    outputs = op(tensor)
                    assert outputs.size(1) == 16 and outputs.size(
                        2) == 32 // stride


def test_darts_series():

    tensor = torch.randn(16, 16, 32, 32)

    # test avg pool bn
    avg_pool_bn = dict(
        type='DartsPoolBN',
        in_channels=16,
        out_channels=16,
        kernel_size=3,
        pool_type='avg',
        stride=1)

    op = OPS.build(avg_pool_bn)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test max pool bn
    max_pool_bn = dict(
        type='DartsPoolBN',
        in_channels=16,
        out_channels=16,
        kernel_size=3,
        pool_type='max',
        stride=1)

    op = OPS.build(max_pool_bn)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test DartsSepConv
    sep_conv = dict(
        type='DartsSepConv',
        in_channels=16,
        out_channels=16,
        kernel_size=3,
        stride=1)

    op = OPS.build(sep_conv)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test DartsSepConv
    sep_conv = dict(
        type='DartsSepConv',
        in_channels=16,
        out_channels=16,
        kernel_size=3,
        stride=1)

    op = OPS.build(sep_conv)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test DartsDilConv
    dil_conv = dict(
        type='DartsDilConv',
        in_channels=16,
        out_channels=16,
        kernel_size=3,
        stride=1)

    op = OPS.build(dil_conv)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

    # test DartsSkipConnect
    skip_connect = dict(
        type='DartsSkipConnect', in_channels=16, out_channels=16, stride=1)

    op = OPS.build(skip_connect)

    # test forward
    outputs = op(tensor)
    assert outputs.size(1) == 16 and outputs.size(2) == 32

tests/test_models/test_op.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
from copy import deepcopy

import pytest
import torch
from mmcv import ConfigDict, digit_version

from mmrazor.models.builder import ARCHITECTURES, PRUNERS


def test_ratio_pruner():
    model_cfg = dict(
        type='mmcls.ImageClassifier',
        backbone=dict(
            type='mmcls.ResNet',
            depth=18,
            num_stages=4,
            out_indices=(3, ),
            style='pytorch'),
        neck=dict(type='mmcls.GlobalAveragePooling'),
        head=dict(
            type='mmcls.LinearClsHead',
            num_classes=1000,
            in_channels=512,
            loss=dict(type='mmcls.CrossEntropyLoss', loss_weight=1.0),
            topk=(1, 5),
        ))

    architecture_cfg = dict(
        type='MMClsArchitecture',
        model=model_cfg,
    )

    pruner_cfg = dict(
        type='RatioPruner',
        ratios=[1 / 8, 2 / 8, 3 / 8, 4 / 8, 5 / 8, 6 / 8, 7 / 8, 1.0])

    _test_reset_bn_running_stats(architecture_cfg, pruner_cfg, False)
    with pytest.raises(AssertionError):
        _test_reset_bn_running_stats(architecture_cfg, pruner_cfg, True)

    imgs = torch.randn(16, 3, 224, 224)
    label = torch.randint(0, 1000, (16, ))

    architecture = ARCHITECTURES.build(architecture_cfg)
    pruner = PRUNERS.build(pruner_cfg)

    pruner.prepare_from_supernet(architecture)
    assert hasattr(pruner, 'channel_spaces')

    # test set_min_channel
    pruner_cfg_ = deepcopy(pruner_cfg)
    pruner_cfg_['ratios'].insert(0, 0)
    pruner_ = PRUNERS.build(pruner_cfg_)
    architecture_ = ARCHITECTURES.build(architecture_cfg)
    pruner_.prepare_from_supernet(architecture_)
    with pytest.raises(AssertionError):
        # Output channels should be a positive integer not zero
        pruner_.set_min_channel()

    # test set_max_channel
    pruner.set_max_channel()
    for name, module in architecture.model.named_modules():
        if hasattr(module, 'in_mask'):
            assert module.in_mask.sum() == module.in_mask.numel()
        if hasattr(module, 'out_mask'):
            assert module.out_mask.sum() == module.out_mask.numel()

    # test channel bins
    pruner.set_min_channel()
    channel_bins_dict = pruner.get_max_channel_bins(max_channel_bins=4)
    pruner.set_channel_bins(channel_bins_dict, 4)
    for name, module in architecture.model.named_modules():
        if hasattr(module, 'in_mask'):
            assert module.in_mask.sum() == module.in_mask.numel()
        if hasattr(module, 'out_mask'):
            assert module.out_mask.sum() == module.out_mask.numel()

    # test making groups logic
    subnet_dict = pruner.sample_subnet()
    assert isinstance(subnet_dict, dict)
    pruner.set_subnet(subnet_dict)
    subnet_dict = pruner.export_subnet()
    assert isinstance(subnet_dict, dict)
    pruner.deploy_subnet(architecture, subnet_dict)
    losses = architecture(imgs, return_loss=True, gt_label=label)
    assert losses['loss'].item() > 0

    # test models with shared module
    model_cfg = ConfigDict(
        type='mmdet.RetinaNet',
        backbone=dict(
            type='ResNet',
            depth=50,
            num_stages=4,
            out_indices=(0, 1, 2, 3),
            frozen_stages=1,
            norm_cfg=dict(type='BN', requires_grad=True),
            norm_eval=True,
            style='pytorch'),
        neck=dict(
            type='FPN',
            in_channels=[256, 512, 1024, 2048],
            out_channels=256,
            start_level=1,
            add_extra_convs='on_input',
            num_outs=5),
        bbox_head=dict(
            type='RetinaHead',
            num_classes=80,
            in_channels=256,
            stacked_convs=4,
            feat_channels=256,
            anchor_generator=dict(
                type='AnchorGenerator',
                octave_base_scale=4,
                scales_per_octave=3,
                ratios=[0.5, 1.0, 2.0],
                strides=[8, 16, 32, 64, 128]),
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[.0, .0, .0, .0],
                target_stds=[1.0, 1.0, 1.0, 1.0]),
            loss_cls=dict(
                type='FocalLoss',
                use_sigmoid=True,
                gamma=2.0,
                alpha=0.25,
                loss_weight=1.0),
            loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
        # model training and testing settings
        train_cfg=dict(
            assigner=dict(
                type='MaxIoUAssigner',
                pos_iou_thr=0.5,
                neg_iou_thr=0.4,
                min_pos_iou=0,
                ignore_iof_thr=-1),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        test_cfg=dict(
            nms_pre=1000,
            min_bbox_size=0,
            score_thr=0.05,
            nms=dict(type='nms', iou_threshold=0.5),
            max_per_img=100))

    architecture_cfg = dict(
        type='MMDetArchitecture',
        model=model_cfg,
    )

    pruner_cfg = dict(
        type='RatioPruner',
        ratios=[1 / 8, 2 / 8, 3 / 8, 4 / 8, 5 / 8, 6 / 8, 7 / 8, 1.0])

    architecture = ARCHITECTURES.build(architecture_cfg)
    pruner = PRUNERS.build(pruner_cfg)
    pruner.prepare_from_supernet(architecture)
    subnet_dict = pruner.sample_subnet()
    pruner.set_subnet(subnet_dict)
    subnet_dict = pruner.export_subnet()
    pruner.deploy_subnet(architecture, subnet_dict)
    architecture.forward_dummy(imgs)

    # test models with concat operations
    model_cfg = ConfigDict(
        type='mmdet.YOLOX',
        input_size=(640, 640),
        random_size_range=(15, 25),
        random_size_interval=10,
        backbone=dict(type='CSPDarknet', deepen_factor=0.33, widen_factor=0.5),
        neck=dict(
            type='YOLOXPAFPN',
            in_channels=[128, 256, 512],
            out_channels=128,
            num_csp_blocks=1),
        bbox_head=dict(
            type='YOLOXHead',
            num_classes=80,
            in_channels=128,
            feat_channels=128),
        train_cfg=dict(
            assigner=dict(type='SimOTAAssigner', center_radius=2.5)),
        # In order to align the source code, the threshold of the val phase is
        # 0.01, and the threshold of the test phase is 0.001.
        test_cfg=dict(
            score_thr=0.01, nms=dict(type='nms', iou_threshold=0.65)))

    architecture_cfg = dict(
        type='MMDetArchitecture',
        model=model_cfg,
    )

    architecture = ARCHITECTURES.build(architecture_cfg)
    pruner.prepare_from_supernet(architecture)
    subnet_dict = pruner.sample_subnet()
    pruner.set_subnet(subnet_dict)
    subnet_dict = pruner.export_subnet()
    pruner.deploy_subnet(architecture, subnet_dict)
    architecture.forward_dummy(imgs)

    # test models with groupnorm
    model_cfg = ConfigDict(
        type='mmdet.ATSS',
        backbone=dict(
            type='ResNet',
            depth=50,
            num_stages=4,
            out_indices=(0, 1, 2, 3),
            frozen_stages=1,
            norm_cfg=dict(type='BN', requires_grad=True),
            norm_eval=True,
            style='pytorch',
            init_cfg=dict(
                type='Pretrained', checkpoint='torchvision://resnet50')),
        neck=dict(
            type='FPN',
            in_channels=[256, 512, 1024, 2048],
            out_channels=256,
            start_level=1,
            add_extra_convs='on_output',
            num_outs=5),
        bbox_head=dict(
            type='ATSSHead',
            num_classes=80,
            in_channels=256,
            stacked_convs=4,
            feat_channels=256,
            anchor_generator=dict(
                type='AnchorGenerator',
                ratios=[1.0],
                octave_base_scale=8,
                scales_per_octave=1,
                strides=[8, 16, 32, 64, 128]),
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[.0, .0, .0, .0],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            loss_cls=dict(
                type='FocalLoss',
                use_sigmoid=True,
                gamma=2.0,
                alpha=0.25,
                loss_weight=1.0),
            loss_bbox=dict(type='GIoULoss', loss_weight=2.0),
            loss_centerness=dict(
                type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0)),
        # training and testing settings
        train_cfg=dict(
            assigner=dict(type='ATSSAssigner', topk=9),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        test_cfg=dict(
            nms_pre=1000,
            min_bbox_size=0,
            score_thr=0.05,
            nms=dict(type='nms', iou_threshold=0.6),
            max_per_img=100))

    architecture_cfg = dict(
        type='MMDetArchitecture',
        model=model_cfg,
    )

    architecture = ARCHITECTURES.build(architecture_cfg)
    # ``StructurePruner`` requires pytorch>=1.6.0 to auto-trace GroupNorm
    # correctly
    min_required_version = '1.6.0'
    if digit_version(torch.__version__) < digit_version(min_required_version):
        with pytest.raises(AssertionError):
            pruner.prepare_from_supernet(architecture)
    else:
        pruner.prepare_from_supernet(architecture)
        subnet_dict = pruner.sample_subnet()
        pruner.set_subnet(subnet_dict)
        subnet_dict = pruner.export_subnet()
        pruner.deploy_subnet(architecture, subnet_dict)
        architecture.forward_dummy(imgs)


def _test_reset_bn_running_stats(architecture_cfg, pruner_cfg, should_fail):
    import os
    import random

    import numpy as np

    def set_seed(seed: int) -> None:
        random.seed(seed)
        os.environ['PYTHONHASHSEED'] = str(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)

    output_list = []

    def output_hook(self, input, output) -> None:
        output_list.append(output)

    set_seed(1024)

    imgs = torch.randn(16, 3, 224, 224)

    torch_rng_state = torch.get_rng_state()
    np_rng_state = np.random.get_state()
    random_rng_state = random.getstate()

    architecture1 = ARCHITECTURES.build(architecture_cfg)
    pruner1 = PRUNERS.build(pruner_cfg)
    if should_fail:
        pruner1._reset_norm_running_stats = lambda *_: None
    set_seed(1)
    pruner1.prepare_from_supernet(architecture1)
    architecture1.model.head.fc.register_forward_hook(output_hook)
    architecture1.eval()
    architecture1(imgs, return_loss=False)

    set_seed(1024)
    torch.set_rng_state(torch_rng_state)
    np.random.set_state(np_rng_state)
    random.setstate(random_rng_state)

    architecture2 = ARCHITECTURES.build(architecture_cfg)
    pruner2 = PRUNERS.build(pruner_cfg)
    if should_fail:
        pruner2._reset_norm_running_stats = lambda *_: None
    set_seed(2)
    pruner2.prepare_from_supernet(architecture2)
    architecture2.model.head.fc.register_forward_hook(output_hook)
    architecture2.eval()
    architecture2(imgs, return_loss=False)

    assert torch.equal(output_list[0].norm(p='fro'),
                       output_list[1].norm(p='fro'))

tests/test_models/test_pruner.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import tempfile

from mmrazor.utils import find_latest_checkpoint


def test_find_latest_checkpoint():
    with tempfile.TemporaryDirectory() as tmpdir:
        path = tmpdir
        latest = find_latest_checkpoint(path)
        # There are no checkpoints in the path.
        assert latest is None

        path = tmpdir + '/none'
        latest = find_latest_checkpoint(path)
        # The path does not exist.
        assert latest is None

    with tempfile.TemporaryDirectory() as tmpdir:
        with open(tmpdir + '/latest.pth', 'w') as f:
            f.write('latest')
        path = tmpdir
        latest = find_latest_checkpoint(path)
        assert latest == osp.join(tmpdir, 'latest.pth')

    with tempfile.TemporaryDirectory() as tmpdir:
        with open(tmpdir + '/iter_4000.pth', 'w') as f:
            f.write('iter_4000')
        with open(tmpdir + '/iter_8000.pth', 'w') as f:
            f.write('iter_8000')
        path = tmpdir
        latest = find_latest_checkpoint(path)
        assert latest == osp.join(tmpdir, 'iter_8000.pth')

    with tempfile.TemporaryDirectory() as tmpdir:
        with open(tmpdir + '/epoch_1.pth', 'w') as f:
            f.write('epoch_1')
        with open(tmpdir + '/epoch_2.pth', 'w') as f:
            f.write('epoch_2')
        path = tmpdir
        latest = find_latest_checkpoint(path)
        assert latest == osp.join(tmpdir, 'epoch_2.pth')

tests/test_utils/test_misc.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import multiprocessing as mp
import os
import platform

import cv2
from mmcv import Config

from mmrazor.utils import setup_multi_processes


def test_setup_multi_processes():
    # temp save system setting
    sys_start_mehod = mp.get_start_method(allow_none=True)
    sys_cv_threads = cv2.getNumThreads()
    # pop and temp save system env vars
    sys_omp_threads = os.environ.pop('OMP_NUM_THREADS', default=None)
    sys_mkl_threads = os.environ.pop('MKL_NUM_THREADS', default=None)

    # test config without setting env
    config = dict(data=dict(workers_per_gpu=2))
    cfg = Config(config)
    setup_multi_processes(cfg)
    assert os.getenv('OMP_NUM_THREADS') == '1'
    assert os.getenv('MKL_NUM_THREADS') == '1'
    # when set to 0, the num threads will be 1
    assert cv2.getNumThreads() == 1
    if platform.system() != 'Windows':
        assert mp.get_start_method() == 'fork'

    # test num workers <= 1
    os.environ.pop('OMP_NUM_THREADS')
    os.environ.pop('MKL_NUM_THREADS')
    config = dict(data=dict(workers_per_gpu=0))
    cfg = Config(config)
    setup_multi_processes(cfg)
    assert 'OMP_NUM_THREADS' not in os.environ
    assert 'MKL_NUM_THREADS' not in os.environ

    # test manually set env var
    os.environ['OMP_NUM_THREADS'] = '4'
    config = dict(data=dict(workers_per_gpu=2))
    cfg = Config(config)
    setup_multi_processes(cfg)
    assert os.getenv('OMP_NUM_THREADS') == '4'

    # test manually set opencv threads and mp start method
    config = dict(
        data=dict(workers_per_gpu=2),
        opencv_num_threads=4,
        mp_start_method='spawn')
    cfg = Config(config)
    setup_multi_processes(cfg)
    assert cv2.getNumThreads() == 4
    assert mp.get_start_method() == 'spawn'

    # revert setting to avoid affecting other programs
    if sys_start_mehod:
        mp.set_start_method(sys_start_mehod, force=True)
    cv2.setNumThreads(sys_cv_threads)
    if sys_omp_threads:
        os.environ['OMP_NUM_THREADS'] = sys_omp_threads
    else:
        os.environ.pop('OMP_NUM_THREADS')
    if sys_mkl_threads:
        os.environ['MKL_NUM_THREADS'] = sys_mkl_threads
    else:
        os.environ.pop('MKL_NUM_THREADS')

tests/test_utils/test_setup_env.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import argparse

import numpy as np
import torch
from mmcv import Config, DictAction
from mmcv.cnn.utils import get_model_complexity_info

from mmrazor.models import build_algorithm


def parse_args():
    parser = argparse.ArgumentParser(description='Get model flops and params')
    parser.add_argument('config', help='config file path')
    parser.add_argument(
        '--shape',
        type=int,
        nargs='+',
        default=[224, 224],
        help='input image size')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--size-divisor',
        type=int,
        default=32,
        help='Pad the input image, the minimum size that is divisible '
        'by size_divisor, -1 means do not pad the image.')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()

    if len(args.shape) == 1:
        h = w = args.shape[0]
    elif len(args.shape) == 2:
        h, w = args.shape
    else:
        raise ValueError('invalid input shape')
    orig_shape = (3, h, w)
    divisor = args.size_divisor
    if divisor > 0:
        h = int(np.ceil(h / divisor)) * divisor
        w = int(np.ceil(w / divisor)) * divisor

    input_shape = (3, h, w)

    cfg = Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    algorithm = build_algorithm(cfg.algorithm)
    if torch.cuda.is_available():
        algorithm.cuda()
    algorithm.eval()

    if hasattr(algorithm.architecture, 'forward_dummy'):
        algorithm.architecture.forward = algorithm.architecture.forward_dummy
    else:
        raise NotImplementedError(
            'FLOPs counter is currently not currently supported with {}'.
            format(algorithm.architecture.__class__.__name__))

    flops, params = get_model_complexity_info(algorithm.architecture,
                                              input_shape)
    split_line = '=' * 30

    if divisor > 0 and \
            input_shape != orig_shape:
        print(f'{split_line}\nUse size divisor set input shape '
              f'from {orig_shape} to {input_shape}\n')
    print(f'{split_line}\nInput shape: {input_shape}\n'
          f'Flops: {flops}\nParams: {params}\n{split_line}')
    print('!!!Please be cautious if you use the results in papers. '
          'You may need to check if all ops are supported and verify that the '
          'flops computation is correct.')


if __name__ == '__main__':
    main()

tools/misc/get_flops.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
from mmcls.apis import multi_gpu_test, single_gpu_test
from mmcls.datasets import build_dataloader, build_dataset
from mmcls.utils import collect_env, get_root_logger
from mmcv import DictAction
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import init_dist, load_checkpoint

from mmrazor.core import build_searcher
from mmrazor.models import build_algorithm
from mmrazor.utils import setup_multi_processes

# TODO import `wrap_fp16_model` from mmcv and delete them from mmcls
try:
    from mmcv.runner import wrap_fp16_model
except ImportError:
    warnings.warn('wrap_fp16_model from mmcls will be deprecated.'
                  'Please install mmcv>=1.1.4.')
    from mmcls.core import wrap_fp16_model


def parse_args():
    parser = argparse.ArgumentParser(
        description='MMClsArchitecture search subnet')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument(
        '--work-dir',
        help='working direction is '
        'to save search result and log')
    parser.add_argument(
        '--resume-from', type=str, help='the checkpoint file to resume from')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument(
        '--device',
        choices=['cpu', 'cuda'],
        default='cuda',
        help='device used for testing')
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)
    return args


def main():
    args = parse_args()

    cfg = mmcv.Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    cfg.model.pretrained = None
    cfg.data.test.test_mode = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)

    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    # build the dataloader
    dataset = build_dataset(cfg.data.test)

    # the extra round_up data will be removed during gpu/cpu collect
    data_loader = build_dataloader(
        dataset,
        samples_per_gpu=cfg.data.samples_per_gpu,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=distributed,
        shuffle=False,
        round_up=True)

    # build the algorithm and load checkpoint
    algorithm = build_algorithm(cfg.algorithm)
    model = algorithm.architecture.model

    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)

    checkpoint = load_checkpoint(
        algorithm, args.checkpoint, map_location='cpu')

    if 'CLASSES' in checkpoint.get('meta', {}):
        CLASSES = checkpoint['meta']['CLASSES']
    else:
        from mmcls.datasets import ImageNet
        warnings.simplefilter('once')
        warnings.warn('Class names are not saved in the checkpoint\'s '
                      'meta data, use imagenet by default.')
        CLASSES = ImageNet.CLASSES

    if not distributed:
        if args.device == 'cpu':
            algorithm = algorithm.cpu()
        else:
            algorithm = MMDataParallel(algorithm, device_ids=[0])
        model.CLASSES = CLASSES
        test_fn = single_gpu_test

    else:
        algorithm = MMDistributedDataParallel(
            algorithm.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False)
        test_fn = multi_gpu_test

    logger.info('build search...')
    searcher = build_searcher(
        cfg.searcher,
        default_args=dict(
            algorithm=algorithm,
            dataloader=data_loader,
            test_fn=test_fn,
            work_dir=cfg.work_dir,
            logger=logger,
            resume_from=args.resume_from))
    logger.info('start search...')
    searcher.search()


if __name__ == '__main__':
    main()

tools/mmcls/search_mmcls.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os
import warnings

import mmcv
import numpy as np
import torch
from mmcls.apis import multi_gpu_test, single_gpu_test
from mmcls.datasets import build_dataloader, build_dataset
from mmcv import DictAction
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import get_dist_info, init_dist, load_checkpoint

from mmrazor.models.builder import build_algorithm
from mmrazor.utils import setup_multi_processes

# TODO import `wrap_fp16_model` from mmcv and delete them from mmcls
try:
    from mmcv.runner import wrap_fp16_model
except ImportError:
    warnings.warn('wrap_fp16_model from mmcls will be deprecated.'
                  'Please install mmcv>=1.1.4.')
    from mmcls.core import wrap_fp16_model


def parse_args():
    parser = argparse.ArgumentParser(description='mmcls test model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument('--out', help='output result file')
    parser.add_argument(
        '--metrics',
        type=str,
        nargs='+',
        help='evaluation metrics, which depends on the dataset, e.g., '
        '"accuracy", "precision", "recall", "f1_score", "support" for single '
        'label dataset, and "mAP", "CP", "CR", "CF1", "OP", "OR", "OF1" for '
        'multi-label dataset')
    parser.add_argument('--show', action='store_true', help='show results')
    parser.add_argument(
        '--show-dir', help='directory where painted images will be saved')
    parser.add_argument(
        '--gpu_collect',
        action='store_true',
        help='whether to use gpu to collect results')
    parser.add_argument('--tmpdir', help='tmp dir for writing some results')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--metric-options',
        nargs='+',
        action=DictAction,
        default={},
        help='custom options for evaluation, the key-value pair in xxx=yyy '
        'format will be parsed as a dict metric_options for dataset.evaluate()'
        ' function.')
    parser.add_argument(
        '--show-options',
        nargs='+',
        action=DictAction,
        help='custom options for show_result. key-value pair in xxx=yyy.'
        'Check available options in `model.show_result`.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument(
        '--device',
        choices=['cpu', 'cuda'],
        default='cuda',
        help='device used for testing')
    parser.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed testing)')
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)
    return args


def main():
    args = parse_args()

    cfg = mmcv.Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    cfg.data.test.test_mode = True

    assert args.metrics or args.out, \
        'Please specify at least one of output path and evaluation metrics.'

    cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)

    # build the dataloader
    dataset = build_dataset(cfg.data.test, default_args=dict(test_mode=True))

    # build the dataloader
    # The default loader config
    loader_cfg = dict(
        # cfg.gpus will be ignored if distributed
        num_gpus=len(cfg.gpu_ids),
        dist=distributed,
        round_up=True,
    )
    # The overall dataloader settings
    loader_cfg.update({
        k: v
        for k, v in cfg.data.items() if k not in [
            'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
            'test_dataloader'
        ]
    })
    test_loader_cfg = {
        **loader_cfg,
        'shuffle': False,  # Not shuffle by default
        'sampler_cfg': None,  # Not use sampler by default
        **cfg.data.get('test_dataloader', {}),
    }
    # the extra round_up data will be removed during gpu/cpu collect
    data_loader = build_dataloader(dataset, **test_loader_cfg)

    # build the algorithm and load checkpoint
    algorithm = build_algorithm(cfg.algorithm)
    model = algorithm.architecture.model
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    checkpoint = load_checkpoint(
        algorithm, args.checkpoint, map_location='cpu')

    if 'CLASSES' in checkpoint.get('meta', {}):
        CLASSES = checkpoint['meta']['CLASSES']
    else:
        from mmcls.datasets import ImageNet
        warnings.simplefilter('once')
        warnings.warn('Class names are not saved in the checkpoint\'s '
                      'meta data, use imagenet by default.')
        CLASSES = ImageNet.CLASSES

    if not distributed:
        if args.device == 'cpu':
            algorithm = algorithm.cpu()
        else:
            algorithm = MMDataParallel(algorithm, device_ids=cfg.gpu_ids)
        model.CLASSES = CLASSES
        show_kwargs = {} if args.show_options is None else args.show_options
        outputs = single_gpu_test(algorithm, data_loader, args.show,
                                  args.show_dir, **show_kwargs)
    else:
        algorithm = MMDistributedDataParallel(
            algorithm.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False)
        outputs = multi_gpu_test(algorithm, data_loader, args.tmpdir,
                                 args.gpu_collect)

    rank, _ = get_dist_info()
    if rank == 0:
        results = {}
        if args.metrics:
            eval_results = dataset.evaluate(outputs, args.metrics,
                                            args.metric_options)
            results.update(eval_results)
            for k, v in eval_results.items():
                print(f'\n{k} : {v:.2f}')
        if args.out:
            scores = np.vstack(outputs)
            pred_score = np.max(scores, axis=1)
            pred_label = np.argmax(scores, axis=1)
            pred_class = [CLASSES[lb] for lb in pred_label]
            results.update({
                'class_scores': scores,
                'pred_score': pred_score,
                'pred_label': pred_label,
                'pred_class': pred_class
            })
            print(f'\ndumping results to {args.out}')
            mmcv.dump(results, args.out)


if __name__ == '__main__':
    main()

tools/mmcls/test_mmcls.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import copy
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
import torch.distributed as dist
from mmcls import __version__
from mmcls.datasets import build_dataset
from mmcls.utils import collect_env, get_root_logger
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist

# Differences from mmclassification
from mmrazor.apis import init_random_seed, set_random_seed, train_mmcls_model
from mmrazor.models import build_algorithm
from mmrazor.utils import setup_multi_processes


def parse_args():
    parser = argparse.ArgumentParser(description='Train a model')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--auto-resume',
        action='store_true',
        help='resume from the latest checkpoint automatically')
    parser.add_argument(
        '--no-validate',
        action='store_true',
        help='whether not to evaluate the checkpoint during training')
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument('--device', help='device used for training')
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='(Deprecated, please use --gpu-id) number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--diff_seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    if args.resume_from is not None:
        cfg.resume_from = args.resume_from
    cfg.auto_resume = args.auto_resume
    if args.gpus is not None:
        cfg.gpu_ids = range(1)
        warnings.warn('`--gpus` is deprecated because we only support '
                      'single GPU mode in non-distributed training. '
                      'Use `gpus=1` now.')
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed training. Use the first GPU '
                      'in `gpu_ids` now.')
    if args.gpus is None and args.gpu_ids is None:
        cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info

    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    # set random seeds
    seed = init_random_seed(args.seed)
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)
    cfg.seed = seed
    meta['seed'] = seed
    meta['exp_name'] = osp.basename(args.config)

    # Difference from mmclassification
    # replace `model` to `algorithm`
    algorithm = build_algorithm(cfg.algorithm)
    algorithm.init_weights()

    datasets = [build_dataset(cfg.data.train)]
    if len(cfg.workflow) == 2:
        val_dataset = copy.deepcopy(cfg.data.val)
        val_dataset.pipeline = cfg.data.train.pipeline
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmcls version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmcls_version=__version__,
            config=cfg.pretty_text,
            CLASSES=datasets[0].CLASSES)
    # add an attribute for visualization convenience
    train_mmcls_model(
        # Difference from mmclassification
        # replace `model` to `algorithm`
        algorithm,
        datasets,
        cfg,
        distributed=distributed,
        validate=(not args.no_validate),
        timestamp=timestamp,
        device='cpu' if args.device == 'cpu' else 'cuda',
        meta=meta)


if __name__ == '__main__':
    main()

tools/mmcls/train_mmcls.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmdet
except (ImportError, ModuleNotFoundError):
    mmdet = None

if mmdet is None:
    raise RuntimeError('mmdet is not installed')
import argparse
import os
import os.path as osp
import time

import mmcv
import torch
from mmcv import Config, DictAction
from mmcv.cnn import fuse_conv_bn
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import init_dist, load_checkpoint, wrap_fp16_model
from mmdet.apis import multi_gpu_test, single_gpu_test
from mmdet.datasets import (build_dataloader, build_dataset,
                            replace_ImageToTensor)
from mmdet.utils import collect_env, get_root_logger

from mmrazor.core import build_searcher
from mmrazor.models import build_algorithm
from mmrazor.utils import setup_multi_processes


def parse_args():
    parser = argparse.ArgumentParser(
        description='MMClsArchitecture search subnet')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument(
        '--work-dir',
        help='working direction is '
        'to save search result and log')
    parser.add_argument('--resume-from', help='resume searching')
    parser.add_argument(
        '--fuse-conv-bn',
        action='store_true',
        help='Whether to fuse conv and bn, this will slightly increase'
        'the inference speed')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # import modules from string list.
    if cfg.get('custom_imports', None):
        from mmcv.utils import import_modules_from_strings
        import_modules_from_strings(**cfg['custom_imports'])
    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    model_cfg = cfg.algorithm.architecture.model
    model_cfg.pretrained = None
    if model_cfg.get('neck'):
        if isinstance(model_cfg.neck, list):
            for neck_cfg in model_cfg.neck:
                if neck_cfg.get('rfp_backbone'):
                    if neck_cfg.rfp_backbone.get('pretrained'):
                        neck_cfg.rfp_backbone.pretrained = None
        elif model_cfg.neck.get('rfp_backbone'):
            if model_cfg.neck.rfp_backbone.get('pretrained'):
                model_cfg.neck.rfp_backbone.pretrained = None

    # in case the test dataset is concatenated
    samples_per_gpu = 1
    if isinstance(cfg.data.test, dict):
        cfg.data.test.test_mode = True
        samples_per_gpu = cfg.data.test.pop('samples_per_gpu', 1)
        if samples_per_gpu > 1:
            # Replace 'ImageToTensor' to 'DefaultFormatBundle'
            cfg.data.test.pipeline = replace_ImageToTensor(
                cfg.data.test.pipeline)
    elif isinstance(cfg.data.test, list):
        for ds_cfg in cfg.data.test:
            ds_cfg.test_mode = True
        samples_per_gpu = max(
            [ds_cfg.pop('samples_per_gpu', 1) for ds_cfg in cfg.data.test])
        if samples_per_gpu > 1:
            for ds_cfg in cfg.data.test:
                ds_cfg.pipeline = replace_ImageToTensor(ds_cfg.pipeline)

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)

    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    # build the dataloader
    dataset = build_dataset(cfg.data.test)
    data_loader = build_dataloader(
        dataset,
        samples_per_gpu=samples_per_gpu,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=distributed,
        shuffle=False)

    # build the model and load checkpoint
    model_cfg.train_cfg = None
    algorithm = build_algorithm(cfg.algorithm)
    model = algorithm.architecture.model

    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    checkpoint = load_checkpoint(
        algorithm, args.checkpoint, map_location='cpu')

    if args.fuse_conv_bn:
        model = fuse_conv_bn(model)
    # old versions did not save class info in checkpoints, this walkaround is
    # for backward compatibility
    if 'CLASSES' in checkpoint.get('meta', {}):
        model.CLASSES = checkpoint['meta']['CLASSES']
    else:
        model.CLASSES = dataset.CLASSES

    if not distributed:
        algorithm = MMDataParallel(algorithm, device_ids=[0])
        test_fn = single_gpu_test
    else:
        algorithm = MMDistributedDataParallel(
            algorithm.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False)
        test_fn = multi_gpu_test

    logger.info('build search...')
    searcher = build_searcher(
        cfg.searcher,
        default_args=dict(
            algorithm=algorithm,
            dataloader=data_loader,
            test_fn=test_fn,
            work_dir=cfg.work_dir,
            logger=logger,
            resume_from=args.resume_from))
    logger.info('start search...')
    searcher.search()


if __name__ == '__main__':
    main()

tools/mmdet/search_mmdet.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmdet
except (ImportError, ModuleNotFoundError):
    mmdet = None

if mmdet is None:
    raise RuntimeError('mmdet is not installed')
import argparse
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
from mmcv import Config, DictAction
from mmcv.cnn import fuse_conv_bn
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
                         wrap_fp16_model)
from mmdet.apis import multi_gpu_test, single_gpu_test
from mmdet.datasets import (build_dataloader, build_dataset,
                            replace_ImageToTensor)

from mmrazor.models.builder import build_algorithm
from mmrazor.utils import setup_multi_processes


def parse_args():
    parser = argparse.ArgumentParser(
        description='MMDet test (and eval) a model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument(
        '--work-dir',
        help='the directory to save the file containing evaluation metrics')
    parser.add_argument('--out', help='output result file in pickle format')
    parser.add_argument(
        '--fuse-conv-bn',
        action='store_true',
        help='Whether to fuse conv and bn, this will slightly increase'
        'the inference speed')
    parser.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed training)')
    parser.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed testing)')
    parser.add_argument(
        '--format-only',
        action='store_true',
        help='Format the output results without perform evaluation. It is'
        'useful when you want to format the result to a specific format and '
        'submit it to the test server')
    parser.add_argument(
        '--eval',
        type=str,
        nargs='+',
        help='evaluation metrics, which depends on the dataset, e.g., "bbox",'
        ' "segm", "proposal" for COCO, and "mAP", "recall" for PASCAL VOC')
    parser.add_argument('--show', action='store_true', help='show results')
    parser.add_argument(
        '--show-dir', help='directory where painted images will be saved')
    parser.add_argument(
        '--show-score-thr',
        type=float,
        default=0.3,
        help='score threshold (default: 0.3)')
    parser.add_argument(
        '--gpu-collect',
        action='store_true',
        help='whether to use gpu to collect results.')
    parser.add_argument(
        '--tmpdir',
        help='tmp directory used for collecting results from multiple '
        'workers, available when gpu-collect is not specified')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--eval-options',
        nargs='+',
        action=DictAction,
        help='custom options for evaluation, the key-value pair in xxx=yyy '
        'format will be kwargs for dataset.evaluate() function')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def main():
    args = parse_args()

    assert args.out or args.eval or args.format_only or args.show \
        or args.show_dir, \
        ('Please specify at least one operation (save/eval/format/show the '
         'results / save the results) with the argument "--out", "--eval"'
         ', "--format-only", "--show" or "--show-dir"')

    if args.eval and args.format_only:
        raise ValueError('--eval and --format_only cannot be both specified')

    if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
        raise ValueError('The output file must be a pkl file.')

    cfg = Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    cfg.algorithm.architecture.model.pretrained = None
    if cfg.algorithm.architecture.model.get('neck'):
        if isinstance(cfg.algorithm.architecture.model.neck, list):
            for neck_cfg in cfg.algorithm.architecture.neck:
                if neck_cfg.get('rfp_backbone'):
                    if neck_cfg.rfp_backbone.get('pretrained'):
                        neck_cfg.rfp_backbone.pretrained = None
        elif cfg.algorithm.architecture.model.neck.get('rfp_backbone'):
            if cfg.algorithm.architecture.model.neck.rfp_backbone.get(
                    'pretrained'):
                cfg.algorithm.architecture.model.neck.rfp_backbone.pretrained = None  # noqa E501

    # in case the test dataset is concatenated
    samples_per_gpu = 1
    if isinstance(cfg.data.test, dict):
        cfg.data.test.test_mode = True
        samples_per_gpu = cfg.data.test.pop('samples_per_gpu', 1)
        if samples_per_gpu > 1:
            # Replace 'ImageToTensor' to 'DefaultFormatBundle'
            cfg.data.test.pipeline = replace_ImageToTensor(
                cfg.data.test.pipeline)
    elif isinstance(cfg.data.test, list):
        for ds_cfg in cfg.data.test:
            ds_cfg.test_mode = True
        samples_per_gpu = max(
            [ds_cfg.pop('samples_per_gpu', 1) for ds_cfg in cfg.data.test])
        if samples_per_gpu > 1:
            for ds_cfg in cfg.data.test:
                ds_cfg.pipeline = replace_ImageToTensor(ds_cfg.pipeline)

    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed testing. Use the first GPU '
                      'in `gpu_ids` now.')
    else:
        cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)

    rank, _ = get_dist_info()
    # allows not to create
    if args.work_dir is not None and rank == 0:
        mmcv.mkdir_or_exist(osp.abspath(args.work_dir))
        timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
        json_file = osp.join(args.work_dir, f'eval_{timestamp}.json')

    # build the dataloader
    dataset = build_dataset(cfg.data.test)
    data_loader = build_dataloader(
        dataset,
        samples_per_gpu=samples_per_gpu,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=distributed,
        shuffle=False)

    # build the algorithm and load checkpoint
    cfg.algorithm.architecture.model.train_cfg = None
    algorithm = build_algorithm(cfg.algorithm)
    model = algorithm.architecture.model

    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    checkpoint = load_checkpoint(
        algorithm, args.checkpoint, map_location='cpu')
    if args.fuse_conv_bn:
        model = fuse_conv_bn(model)
    # old versions did not save class info in checkpoints, this walkaround is
    # for backward compatibility
    if 'CLASSES' in checkpoint.get('meta', {}):
        model.CLASSES = checkpoint['meta']['CLASSES']
    else:
        model.CLASSES = dataset.CLASSES

    if not distributed:
        algorithm = MMDataParallel(algorithm, device_ids=cfg.gpu_ids)
        outputs = single_gpu_test(algorithm, data_loader, args.show,
                                  args.show_dir, args.show_score_thr)
    else:
        algorithm = MMDistributedDataParallel(
            algorithm.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False)
        outputs = multi_gpu_test(algorithm, data_loader, args.tmpdir,
                                 args.gpu_collect)

    rank, _ = get_dist_info()
    if rank == 0:
        if args.out:
            print(f'\nwriting results to {args.out}')
            mmcv.dump(outputs, args.out)
        kwargs = {} if args.eval_options is None else args.eval_options
        if args.format_only:
            dataset.format_results(outputs, **kwargs)
        if args.eval:
            eval_kwargs = cfg.get('evaluation', {}).copy()
            # hard-code way to remove EvalHook args
            for key in [
                    'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best',
                    'rule', 'dynamic_intervals'
            ]:
                eval_kwargs.pop(key, None)
            eval_kwargs.update(dict(metric=args.eval, **kwargs))
            metric = dataset.evaluate(outputs, **eval_kwargs)
            print(metric)
            metric_dict = dict(config=args.config, metric=metric)
            if args.work_dir is not None and rank == 0:
                mmcv.dump(metric_dict, json_file)


if __name__ == '__main__':
    main()

tools/mmdet/test_mmdet.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmdet
except (ImportError, ModuleNotFoundError):
    mmdet = None

if mmdet is None:
    raise RuntimeError('mmdet is not installed')

import argparse
import copy
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
import torch.distributed as dist
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist
from mmcv.utils import get_git_hash
from mmdet import __version__
from mmdet.datasets import build_dataset
from mmdet.utils import collect_env, get_root_logger

from mmrazor.apis import init_random_seed, set_random_seed, train_mmdet_model
from mmrazor.models import build_algorithm
from mmrazor.utils import setup_multi_processes


def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--auto-resume',
        action='store_true',
        help='resume from the latest checkpoint automatically')
    parser.add_argument(
        '--no-validate',
        action='store_true',
        help='whether not to evaluate the checkpoint during training')
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='(Deprecated, please use --gpu-id) number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--diff_seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # import modules from string list.
    if cfg.get('custom_imports', None):
        from mmcv.utils import import_modules_from_strings
        import_modules_from_strings(**cfg['custom_imports'])
    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    if args.resume_from is not None:
        cfg.resume_from = args.resume_from
    cfg.auto_resume = args.auto_resume
    if args.gpus is not None:
        cfg.gpu_ids = range(1)
        warnings.warn('`--gpus` is deprecated because we only support '
                      'single GPU mode in non-distributed training. '
                      'Use `gpus=1` now.')
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed training. Use the first GPU '
                      'in `gpu_ids` now.')
    if args.gpus is None and args.gpu_ids is None:
        cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        # re-set gpu_ids with distributed training mode
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info
    meta['config'] = cfg.pretty_text
    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    # set random seeds
    seed = init_random_seed(args.seed)
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)
    cfg.seed = seed
    meta['seed'] = seed
    meta['exp_name'] = osp.basename(args.config)

    algorithm = build_algorithm(cfg.algorithm)
    algorithm.init_weights()

    datasets = [build_dataset(cfg.data.train)]
    if len(cfg.workflow) == 2:
        val_dataset = copy.deepcopy(cfg.data.val)
        val_dataset.pipeline = cfg.data.train.pipeline
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmdet version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmdet_version=__version__ + get_git_hash()[:7],
            CLASSES=datasets[0].CLASSES)
    # add an attribute for visualization convenience
    algorithm.CLASSES = datasets[0].CLASSES
    train_mmdet_model(
        algorithm,
        datasets,
        cfg,
        distributed=distributed,
        validate=(not args.no_validate),
        timestamp=timestamp,
        meta=meta)


if __name__ == '__main__':
    main()

tools/mmdet/train_mmdet.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmseg
except (ImportError, ModuleNotFoundError):
    mmseg = None

if mmseg is None:
    raise RuntimeError('mmseg is not installed')

import argparse
import os
import os.path as osp
import shutil
import time
import warnings

import mmcv
import torch
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
                         wrap_fp16_model)
from mmcv.utils import DictAction
from mmseg.apis import multi_gpu_test, single_gpu_test
from mmseg.datasets import build_dataloader, build_dataset

from mmrazor.models.builder import build_algorithm
from mmrazor.utils import setup_multi_processes


def parse_args():
    parser = argparse.ArgumentParser(
        description='mmseg test (and eval) a model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument(
        '--work-dir',
        help=('if specified, the evaluation metric results will be dumped'
              'into the directory as json'))
    parser.add_argument(
        '--aug-test', action='store_true', help='Use Flip and Multi scale aug')
    parser.add_argument('--out', help='output result file in pickle format')
    parser.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed testing)')
    parser.add_argument(
        '--format-only',
        action='store_true',
        help='Format the output results without perform evaluation. It is'
        'useful when you want to format the result to a specific format and '
        'submit it to the test server')
    parser.add_argument(
        '--eval',
        type=str,
        nargs='+',
        help='evaluation metrics, which depends on the dataset, e.g., "mIoU"'
        ' for generic datasets, and "cityscapes" for Cityscapes')
    parser.add_argument('--show', action='store_true', help='show results')
    parser.add_argument(
        '--show-dir', help='directory where painted images will be saved')
    parser.add_argument(
        '--gpu-collect',
        action='store_true',
        help='whether to use gpu to collect results.')
    parser.add_argument(
        '--tmpdir',
        help='tmp directory used for collecting results from multiple '
        'workers, available when gpu_collect is not specified')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--eval-options',
        nargs='+',
        action=DictAction,
        help='custom options for evaluation')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument(
        '--opacity',
        type=float,
        default=0.5,
        help='Opacity of painted segmentation map. In (0, 1] range.')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)
    return args


def main():
    args = parse_args()

    assert args.out or args.eval or args.format_only or args.show \
        or args.show_dir, \
        ('Please specify at least one operation (save/eval/format/show the '
         'results / save the results) with the argument "--out", "--eval"'
         ', "--format-only", "--show" or "--show-dir"')

    if args.eval and args.format_only:
        raise ValueError('--eval and --format_only cannot be both specified')

    if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
        raise ValueError('The output file must be a pkl file.')

    cfg = mmcv.Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    if args.aug_test:
        # hard code index
        cfg.data.test.pipeline[1].img_ratios = [
            0.5, 0.75, 1.0, 1.25, 1.5, 1.75
        ]
        cfg.data.test.pipeline[1].flip = True
    # Difference from mmsegmentation
    cfg.algorithm.architecture.model.pretrained = None
    cfg.data.test.test_mode = True

    cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)

    rank, _ = get_dist_info()
    # allows not to create
    if args.work_dir is not None and rank == 0:
        mmcv.mkdir_or_exist(osp.abspath(args.work_dir))
        timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
        json_file = osp.join(args.work_dir, f'eval_{timestamp}.json')

    # build the dataloader
    # TODO: support multiple images per gpu (only minor changes are needed)
    dataset = build_dataset(cfg.data.test)
    # The default loader config
    loader_cfg = dict(
        # cfg.gpus will be ignored if distributed
        num_gpus=len(cfg.gpu_ids),
        dist=distributed,
        shuffle=False)
    # The overall dataloader settings
    loader_cfg.update({
        k: v
        for k, v in cfg.data.items() if k not in [
            'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
            'test_dataloader'
        ]
    })
    test_loader_cfg = {
        **loader_cfg,
        'samples_per_gpu': 1,
        'shuffle': False,  # Not shuffle by default
        **cfg.data.get('test_dataloader', {})
    }
    # build the dataloader
    data_loader = build_dataloader(dataset, **test_loader_cfg)

    # build the algorithm and load checkpoint
    # Difference from mmsegmentation
    cfg.algorithm.architecture.model.train_cfg = None
    algorithm = build_algorithm(cfg.algorithm)
    model = algorithm.architecture.model

    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    checkpoint = load_checkpoint(
        algorithm, args.checkpoint, map_location='cpu')
    if 'CLASSES' in checkpoint.get('meta', {}):
        model.CLASSES = checkpoint['meta']['CLASSES']
    else:
        print('"CLASSES" not found in meta, use dataset.CLASSES instead')
        model.CLASSES = dataset.CLASSES
    if 'PALETTE' in checkpoint.get('meta', {}):
        model.PALETTE = checkpoint['meta']['PALETTE']
    else:
        print('"PALETTE" not found in meta, use dataset.PALETTE instead')
        model.PALETTE = dataset.PALETTE

    # clean gpu memory when starting a new evaluation.
    torch.cuda.empty_cache()
    eval_kwargs = {} if args.eval_options is None else args.eval_options

    # Deprecated
    efficient_test = eval_kwargs.get('efficient_test', False)
    if efficient_test:
        warnings.warn(
            '``efficient_test=True`` does not have effect in tools/test.py, '
            'the evaluation and format results are CPU memory efficient by '
            'default')

    eval_on_format_results = (
        args.eval is not None and 'cityscapes' in args.eval)
    if eval_on_format_results:
        assert len(args.eval) == 1, 'eval on format results is not ' \
                                    'applicable for metrics other than ' \
                                    'cityscapes'
    if args.format_only or eval_on_format_results:
        if 'imgfile_prefix' in eval_kwargs:
            tmpdir = eval_kwargs['imgfile_prefix']
        else:
            tmpdir = '.format_cityscapes'
            eval_kwargs.setdefault('imgfile_prefix', tmpdir)
        mmcv.mkdir_or_exist(tmpdir)
    else:
        tmpdir = None

    if not distributed:
        algorithm = MMDataParallel(algorithm, device_ids=cfg.gpu_ids)
        results = single_gpu_test(
            algorithm,
            data_loader,
            args.show,
            args.show_dir,
            False,
            args.opacity,
            pre_eval=args.eval is not None and not eval_on_format_results,
            format_only=args.format_only or eval_on_format_results,
            format_args=eval_kwargs)
    else:
        algorithm = MMDistributedDataParallel(
            algorithm.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False)
        results = multi_gpu_test(
            algorithm,
            data_loader,
            args.tmpdir,
            args.gpu_collect,
            False,
            pre_eval=args.eval is not None and not eval_on_format_results,
            format_only=args.format_only or eval_on_format_results,
            format_args=eval_kwargs)

    rank, _ = get_dist_info()
    if rank == 0:
        if args.out:
            warnings.warn(
                'The behavior of ``args.out`` has been changed since MMSeg '
                'v0.16, the pickled outputs could be seg map as type of '
                'np.array, pre-eval results or file paths for '
                '``dataset.format_results()``.')
            print(f'\nwriting results to {args.out}')
            mmcv.dump(results, args.out)
        if args.eval:
            eval_kwargs.update(metric=args.eval)
            metric = dataset.evaluate(results, **eval_kwargs)
            metric_dict = dict(config=args.config, metric=metric)
            if args.work_dir is not None and rank == 0:
                mmcv.dump(metric_dict, json_file, indent=4)
            if tmpdir is not None and eval_on_format_results:
                # remove tmp dir when cityscapes evaluation
                shutil.rmtree(tmpdir)


if __name__ == '__main__':
    main()

tools/mmseg/test_mmseg.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
try:
    import mmseg
except (ImportError, ModuleNotFoundError):
    mmseg = None

if mmseg is None:
    raise RuntimeError('mmseg is not installed')

import argparse
import copy
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
import torch.distributed as dist
from mmcv.cnn.utils import revert_sync_batchnorm
from mmcv.runner import get_dist_info, init_dist
from mmcv.utils import Config, DictAction, get_git_hash
from mmseg import __version__
from mmseg.datasets import build_dataset
from mmseg.utils import collect_env, get_root_logger

# Differences from mmdetection
from mmrazor.apis import init_random_seed, set_random_seed, train_mmseg_model
from mmrazor.models.builder import build_algorithm
from mmrazor.utils import setup_multi_processes


def parse_args():
    parser = argparse.ArgumentParser(description='Train a segmentor')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument(
        '--load-from', help='the checkpoint file to load weights from')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--auto-resume',
        action='store_true',
        help='resume from the latest checkpoint automatically')
    parser.add_argument(
        '--no-validate',
        action='store_true',
        help='whether not to evaluate the checkpoint during training')
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='(Deprecated, please use --gpu-id) number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--diff_seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)

    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    if args.load_from is not None:
        cfg.load_from = args.load_from
    if args.resume_from is not None:
        cfg.resume_from = args.resume_from
    cfg.auto_resume = args.auto_resume
    if args.gpus is not None:
        cfg.gpu_ids = range(1)
        warnings.warn('`--gpus` is deprecated because we only support '
                      'single GPU mode in non-distributed training. '
                      'Use `gpus=1` now.')
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed training. Use the first GPU '
                      'in `gpu_ids` now.')
    if args.gpus is None and args.gpu_ids is None:
        cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        # gpu_ids is used to calculate iter when resuming checkpoint
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info

    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    # set random seeds
    seed = init_random_seed(args.seed)
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {args.seed}, deterministic: '
                f'{args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)
    cfg.seed = seed
    meta['seed'] = args.seed
    meta['exp_name'] = osp.basename(args.config)

    # Difference from mmsegmentation
    # replace `model` to `algorithm`
    algorithm = build_algorithm(cfg.algorithm)
    algorithm.init_weights()

    # SyncBN is not support for DP
    if not distributed:
        warnings.warn(
            'SyncBN is only supported with DDP. To be compatible with DP, '
            'we convert SyncBN to BN. Please use dist_train.sh which can '
            'avoid this error.')
        algorithm = revert_sync_batchnorm(algorithm)

    logger.info(algorithm)

    datasets = [build_dataset(cfg.data.train)]
    if len(cfg.workflow) == 2:
        val_dataset = copy.deepcopy(cfg.data.val)
        val_dataset.pipeline = cfg.data.train.pipeline
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmseg version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmseg_version=f'{__version__}+{get_git_hash()[:7]}',
            config=cfg.pretty_text,
            CLASSES=datasets[0].CLASSES,
            PALETTE=datasets[0].PALETTE)
    # add an attribute for visualization convenience
    # Difference from mmsegmentation
    # replace `model` to `algorithm`
    algorithm.CLASSES = datasets[0].CLASSES
    # passing checkpoint meta for saving best checkpoint
    meta.update(cfg.checkpoint_config.meta)
    train_mmseg_model(
        # Difference from mmsegmentation
        # replace `model` to `algorithm`
        algorithm,
        datasets,
        cfg,
        distributed=distributed,
        validate=(not args.no_validate),
        timestamp=timestamp,
        meta=meta)


if __name__ == '__main__':
    main()

tools/mmseg/train_mmseg.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import datetime
from pathlib import Path

import mmcv
import torch
from mmcv import digit_version


def parse_args():
    parser = argparse.ArgumentParser(
        description='Process a checkpoint to be published')
    parser.add_argument('in_file', help='input checkpoint filename')
    parser.add_argument('out_file', help='output checkpoint filename')
    parser.add_argument('--mutable-cfg', help='input mutable cfg filename')
    parser.add_argument('--channel-cfg', help='output channel cfg filename')
    args = parser.parse_args()
    return args


def cal_file_sha256(file_path: str) -> str:
    import hashlib

    BLOCKSIZE = 65536
    sha256_hash = hashlib.sha256()

    with open(file_path, 'rb') as f:
        block = f.read(BLOCKSIZE)
        while block:
            sha256_hash.update(block)
            block = f.read(BLOCKSIZE)

    return sha256_hash.hexdigest()


def process_checkpoint(in_file,
                       out_file,
                       mutable_cfg_file=None,
                       channel_cfg_file=None):
    checkpoint = torch.load(in_file, map_location='cpu')
    # remove optimizer for smaller file size
    if 'optimizer' in checkpoint:
        del checkpoint['optimizer']
    # if it is necessary to remove some sensitive data in checkpoint['meta'],
    # add the code here.
    if digit_version(torch.__version__) >= digit_version('1.6'):
        torch.save(checkpoint, out_file, _use_new_zipfile_serialization=False)
    else:
        torch.save(checkpoint, out_file)

    sha = cal_file_sha256(out_file)
    if out_file.endswith('.pth'):
        out_file_name = out_file[:-4]
    else:
        out_file_name = out_file

    current_date = datetime.datetime.now().strftime('%Y%m%d')
    final_file_prefix = out_file_name + f'_{current_date}-{sha[:8]}'
    final_ckpt_file = f'{final_file_prefix}.pth'
    Path(out_file).rename(final_ckpt_file)

    print(f'Successfully generated the publish-ckpt as {final_ckpt_file}.')

    if mutable_cfg_file:
        mutable_cfg = mmcv.fileio.load(mutable_cfg_file)
        final_mutable_cfg_file = f'{final_file_prefix}_mutable_cfg.yaml'
        mmcv.fileio.dump(mutable_cfg, final_mutable_cfg_file)
        print(f'Successfully generated the publish-mutable-cfg as \
                {final_mutable_cfg_file}.')

    if channel_cfg_file:
        channel_cfg = mmcv.fileio.load(channel_cfg_file)
        final_channel_cfg_file = f'{final_file_prefix}_channel_cfg.yaml'
        mmcv.fileio.dump(channel_cfg, final_channel_cfg_file)
        print(f'Successfully generated the publish-channel-cfg as \
                {final_channel_cfg_file}.')


def main():
    args = parse_args()
    out_dir = Path(args.out_file).parent
    if not out_dir.exists():
        raise ValueError(f'Directory {out_dir} does not exist, '
                         'please generate it manually.')
    process_checkpoint(args.in_file, args.out_file, args.mutable_cfg,
                       args.channel_cfg)


if __name__ == '__main__':
    main()

tools/model_converters/publish_model.py
Gumpest/AvatarKD
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os.path as osp

from mmcv import Config
from mmcv.runner import load_checkpoint, save_checkpoint

from mmrazor.models import build_algorithm
from mmrazor.models.pruners.utils import SwitchableBatchNorm2d


def parse_args():
    parser = argparse.ArgumentParser(description='Split a slimmable trained'
                                     'model checkpoint')
    parser.add_argument('config', type=str, help='path of train config file')
    parser.add_argument('checkpoint', type=str, help='checkpoint path')
    parser.add_argument(
        '--channel-cfgs',
        nargs='+',
        help='The path of the channel configs. '
        'The order should be the same as that of train.')
    parser.add_argument('--output-dir', type=str, default='')
    args = parser.parse_args()

    return args


def convert_bn(module, bn_ind):

    def traverse(module):
        for name, child in module.named_children():
            if isinstance(child, SwitchableBatchNorm2d):
                setattr(module, name, child.bns[bn_ind])
            else:
                traverse(child)

    traverse(module)


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    cfg.merge_from_dict(dict(algorithm=dict(channel_cfg=args.channel_cfgs)))

    for i, channel_cfg in enumerate(args.channel_cfgs):
        algorithm = build_algorithm(cfg.algorithm)
        load_checkpoint(algorithm, args.checkpoint, map_location='cpu')
        convert_bn(algorithm, i)
        for module in algorithm.modules():
            if hasattr(module, 'out_mask'):
                del module.out_mask
            if hasattr(module, 'in_mask'):
                del module.in_mask
        assert algorithm.with_pruner, \
            'The algorithm should has attr pruner. Please check your ' \
            'config file.'
        algorithm.pruner.deploy_subnet(algorithm.architecture,
                                       algorithm.channel_cfg[i])
        filename = osp.join(args.output_dir, f'checkpoint_{i + 1}.pth')
        save_checkpoint(algorithm, filename)

    print(f'Successfully split the original checkpoint `{args.checkpoint}` to '
          f'{len(args.channel_cfgs)} different checkpoints.')


if __name__ == '__main__':
    main()

tools/model_converters/split_checkpoint.py
Gumpest/AvatarKD