Add official task teachers and align TAIC train/eval configs.

Wire official backbone/neck teachers, freeze_base_codec support, and task-network configs so training and metric eval match the paper setup.

checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py

@@ -0,0 +1,778 @@
+auto_scale_lr = dict(base_batch_size=16, enable=False)
+backbone_embed_multi = dict(decay_mult=0.0, lr_mult=0.1)
+backbone_norm_multi = dict(decay_mult=0.0, lr_mult=0.1)
+backend_args = None
+batch_augments = [
+    dict(
+        img_pad_value=0,
+        mask_pad_value=0,
+        pad_mask=True,
+        pad_seg=True,
+        seg_pad_value=255,
+        size=(
+            1024,
+            1024,
+        ),
+        type='BatchFixedSizePad'),
+]
+custom_keys = dict({
+    'absolute_pos_embed':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone':
+    dict(decay_mult=1.0, lr_mult=0.1),
+    'backbone.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.patch_embed.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.0.blocks.0.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.0.blocks.1.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.0.downsample.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.1.blocks.0.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.1.blocks.1.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.1.downsample.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.0.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.1.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.10.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.11.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.12.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.13.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.14.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.15.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.16.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.17.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.2.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.3.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.4.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.5.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.6.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.7.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.8.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.blocks.9.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.2.downsample.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.3.blocks.0.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'backbone.stages.3.blocks.1.norm':
+    dict(decay_mult=0.0, lr_mult=0.1),
+    'level_embed':
+    dict(decay_mult=0.0, lr_mult=1.0),
+    'query_embed':
+    dict(decay_mult=0.0, lr_mult=1.0),
+    'query_feat':
+    dict(decay_mult=0.0, lr_mult=1.0),
+    'relative_position_bias_table':
+    dict(decay_mult=0.0, lr_mult=0.1)
+})
+data_preprocessor = dict(
+    batch_augments=[
+        dict(
+            img_pad_value=0,
+            mask_pad_value=0,
+            pad_mask=True,
+            pad_seg=True,
+            seg_pad_value=255,
+            size=(
+                1024,
+                1024,
+            ),
+            type='BatchFixedSizePad'),
+    ],
+    bgr_to_rgb=True,
+    mask_pad_value=0,
+    mean=[
+        123.675,
+        116.28,
+        103.53,
+    ],
+    pad_mask=True,
+    pad_seg=True,
+    pad_size_divisor=32,
+    seg_pad_value=255,
+    std=[
+        58.395,
+        57.12,
+        57.375,
+    ],
+    type='DetDataPreprocessor')
+data_root = 'data/coco/'
+dataset_type = 'CocoPanopticDataset'
+default_hooks = dict(
+    checkpoint=dict(
+        by_epoch=False,
+        interval=5000,
+        max_keep_ckpts=3,
+        save_last=True,
+        type='CheckpointHook'),
+    logger=dict(interval=50, type='LoggerHook'),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+    timer=dict(type='IterTimerHook'),
+    visualization=dict(type='DetVisualizationHook'))
+default_scope = 'mmdet'
+depths = [
+    2,
+    2,
+    18,
+    2,
+]
+dynamic_intervals = [
+    (
+        365001,
+        368750,
+    ),
+]
+embed_multi = dict(decay_mult=0.0, lr_mult=1.0)
+env_cfg = dict(
+    cudnn_benchmark=False,
+    dist_cfg=dict(backend='nccl'),
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0))
+image_size = (
+    1024,
+    1024,
+)
+interval = 5000
+load_from = None
+log_level = 'INFO'
+log_processor = dict(by_epoch=False, type='LogProcessor', window_size=50)
+max_iters = 368750
+model = dict(
+    backbone=dict(
+        attn_drop_rate=0.0,
+        convert_weights=True,
+        depths=[
+            2,
+            2,
+            18,
+            2,
+        ],
+        drop_path_rate=0.3,
+        drop_rate=0.0,
+        embed_dims=128,
+        frozen_stages=-1,
+        init_cfg=dict(
+            checkpoint=
+            'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth',
+            type='Pretrained'),
+        mlp_ratio=4,
+        num_heads=[
+            4,
+            8,
+            16,
+            32,
+        ],
+        out_indices=(
+            0,
+            1,
+            2,
+            3,
+        ),
+        patch_norm=True,
+        pretrain_img_size=384,
+        qk_scale=None,
+        qkv_bias=True,
+        type='SwinTransformer',
+        window_size=12,
+        with_cp=False),
+    data_preprocessor=dict(
+        batch_augments=[
+            dict(
+                img_pad_value=0,
+                mask_pad_value=0,
+                pad_mask=True,
+                pad_seg=True,
+                seg_pad_value=255,
+                size=(
+                    1024,
+                    1024,
+                ),
+                type='BatchFixedSizePad'),
+        ],
+        bgr_to_rgb=True,
+        mask_pad_value=0,
+        mean=[
+            123.675,
+            116.28,
+            103.53,
+        ],
+        pad_mask=True,
+        pad_seg=True,
+        pad_size_divisor=32,
+        seg_pad_value=255,
+        std=[
+            58.395,
+            57.12,
+            57.375,
+        ],
+        type='DetDataPreprocessor'),
+    init_cfg=None,
+    panoptic_fusion_head=dict(
+        init_cfg=None,
+        loss_panoptic=None,
+        num_stuff_classes=53,
+        num_things_classes=80,
+        type='MaskFormerFusionHead'),
+    panoptic_head=dict(
+        enforce_decoder_input_project=False,
+        feat_channels=256,
+        in_channels=[
+            128,
+            256,
+            512,
+            1024,
+        ],
+        loss_cls=dict(
+            class_weight=[
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                1.0,
+                0.1,
+            ],
+            loss_weight=2.0,
+            reduction='mean',
+            type='CrossEntropyLoss',
+            use_sigmoid=False),
+        loss_dice=dict(
+            activate=True,
+            eps=1.0,
+            loss_weight=5.0,
+            naive_dice=True,
+            reduction='mean',
+            type='DiceLoss',
+            use_sigmoid=True),
+        loss_mask=dict(
+            loss_weight=5.0,
+            reduction='mean',
+            type='CrossEntropyLoss',
+            use_sigmoid=True),
+        num_queries=100,
+        num_stuff_classes=53,
+        num_things_classes=80,
+        num_transformer_feat_level=3,
+        out_channels=256,
+        pixel_decoder=dict(
+            act_cfg=dict(type='ReLU'),
+            encoder=dict(
+                layer_cfg=dict(
+                    ffn_cfg=dict(
+                        act_cfg=dict(inplace=True, type='ReLU'),
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        ffn_drop=0.0,
+                        num_fcs=2),
+                    self_attn_cfg=dict(
+                        batch_first=True,
+                        dropout=0.0,
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4)),
+                num_layers=6),
+            norm_cfg=dict(num_groups=32, type='GN'),
+            num_outs=3,
+            positional_encoding=dict(normalize=True, num_feats=128),
+            type='MSDeformAttnPixelDecoder'),
+        positional_encoding=dict(normalize=True, num_feats=128),
+        strides=[
+            4,
+            8,
+            16,
+            32,
+        ],
+        transformer_decoder=dict(
+            init_cfg=None,
+            layer_cfg=dict(
+                cross_attn_cfg=dict(
+                    batch_first=True, dropout=0.0, embed_dims=256,
+                    num_heads=8),
+                ffn_cfg=dict(
+                    act_cfg=dict(inplace=True, type='ReLU'),
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    ffn_drop=0.0,
+                    num_fcs=2),
+                self_attn_cfg=dict(
+                    batch_first=True, dropout=0.0, embed_dims=256,
+                    num_heads=8)),
+            num_layers=9,
+            return_intermediate=True),
+        type='Mask2FormerHead'),
+    test_cfg=dict(
+        filter_low_score=True,
+        instance_on=True,
+        iou_thr=0.8,
+        max_per_image=100,
+        panoptic_on=True,
+        semantic_on=False),
+    train_cfg=dict(
+        assigner=dict(
+            match_costs=[
+                dict(type='ClassificationCost', weight=2.0),
+                dict(
+                    type='CrossEntropyLossCost', use_sigmoid=True, weight=5.0),
+                dict(eps=1.0, pred_act=True, type='DiceCost', weight=5.0),
+            ],
+            type='HungarianAssigner'),
+        importance_sample_ratio=0.75,
+        num_points=12544,
+        oversample_ratio=3.0,
+        sampler=dict(type='MaskPseudoSampler')),
+    type='Mask2Former')
+num_classes = 133
+num_stuff_classes = 53
+num_things_classes = 80
+optim_wrapper = dict(
+    clip_grad=dict(max_norm=0.01, norm_type=2),
+    optimizer=dict(
+        betas=(
+            0.9,
+            0.999,
+        ),
+        eps=1e-08,
+        lr=0.0001,
+        type='AdamW',
+        weight_decay=0.05),
+    paramwise_cfg=dict(
+        custom_keys=dict({
+            'absolute_pos_embed':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone':
+            dict(decay_mult=1.0, lr_mult=0.1),
+            'backbone.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.patch_embed.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.0.blocks.0.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.0.blocks.1.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.0.downsample.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.1.blocks.0.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.1.blocks.1.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.1.downsample.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.0.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.1.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.10.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.11.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.12.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.13.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.14.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.15.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.16.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.17.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.2.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.3.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.4.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.5.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.6.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.7.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.8.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.blocks.9.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.2.downsample.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.3.blocks.0.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'backbone.stages.3.blocks.1.norm':
+            dict(decay_mult=0.0, lr_mult=0.1),
+            'level_embed':
+            dict(decay_mult=0.0, lr_mult=1.0),
+            'query_embed':
+            dict(decay_mult=0.0, lr_mult=1.0),
+            'query_feat':
+            dict(decay_mult=0.0, lr_mult=1.0),
+            'relative_position_bias_table':
+            dict(decay_mult=0.0, lr_mult=0.1)
+        }),
+        norm_decay_mult=0.0),
+    type='OptimWrapper')
+param_scheduler = dict(
+    begin=0,
+    by_epoch=False,
+    end=368750,
+    gamma=0.1,
+    milestones=[
+        327778,
+        355092,
+    ],
+    type='MultiStepLR')
+pretrained = 'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth'
+resume = False
+test_cfg = dict(type='TestLoop')
+test_dataloader = dict(
+    batch_size=1,
+    dataset=dict(
+        ann_file='annotations/panoptic_val2017.json',
+        backend_args=None,
+        data_prefix=dict(img='val2017/', seg='annotations/panoptic_val2017/'),
+        data_root='data/coco/',
+        pipeline=[
+            dict(backend_args=None, type='LoadImageFromFile'),
+            dict(keep_ratio=True, scale=(
+                1333,
+                800,
+            ), type='Resize'),
+            dict(backend_args=None, type='LoadPanopticAnnotations'),
+            dict(
+                meta_keys=(
+                    'img_id',
+                    'img_path',
+                    'ori_shape',
+                    'img_shape',
+                    'scale_factor',
+                ),
+                type='PackDetInputs'),
+        ],
+        test_mode=True,
+        type='CocoPanopticDataset'),
+    drop_last=False,
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(shuffle=False, type='DefaultSampler'))
+test_evaluator = [
+    dict(
+        ann_file='data/coco/annotations/panoptic_val2017.json',
+        backend_args=None,
+        seg_prefix='data/coco/annotations/panoptic_val2017/',
+        type='CocoPanopticMetric'),
+    dict(
+        ann_file='data/coco/annotations/instances_val2017.json',
+        backend_args=None,
+        metric=[
+            'bbox',
+            'segm',
+        ],
+        type='CocoMetric'),
+]
+test_pipeline = [
+    dict(backend_args=None, type='LoadImageFromFile'),
+    dict(keep_ratio=True, scale=(
+        1333,
+        800,
+    ), type='Resize'),
+    dict(backend_args=None, type='LoadPanopticAnnotations'),
+    dict(
+        meta_keys=(
+            'img_id',
+            'img_path',
+            'ori_shape',
+            'img_shape',
+            'scale_factor',
+        ),
+        type='PackDetInputs'),
+]
+train_cfg = dict(
+    dynamic_intervals=[
+        (
+            365001,
+            368750,
+        ),
+    ],
+    max_iters=368750,
+    type='IterBasedTrainLoop',
+    val_interval=5000)
+train_dataloader = dict(
+    batch_sampler=dict(type='AspectRatioBatchSampler'),
+    batch_size=2,
+    dataset=dict(
+        ann_file='annotations/panoptic_train2017.json',
+        backend_args=None,
+        data_prefix=dict(
+            img='train2017/', seg='annotations/panoptic_train2017/'),
+        data_root='data/coco/',
+        filter_cfg=dict(filter_empty_gt=True, min_size=32),
+        pipeline=[
+            dict(backend_args=None, to_float32=True, type='LoadImageFromFile'),
+            dict(
+                backend_args=None,
+                type='LoadPanopticAnnotations',
+                with_bbox=True,
+                with_mask=True,
+                with_seg=True),
+            dict(prob=0.5, type='RandomFlip'),
+            dict(
+                keep_ratio=True,
+                ratio_range=(
+                    0.1,
+                    2.0,
+                ),
+                scale=(
+                    1024,
+                    1024,
+                ),
+                type='RandomResize'),
+            dict(
+                allow_negative_crop=True,
+                crop_size=(
+                    1024,
+                    1024,
+                ),
+                crop_type='absolute',
+                recompute_bbox=True,
+                type='RandomCrop'),
+            dict(type='PackDetInputs'),
+        ],
+        type='CocoPanopticDataset'),
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(shuffle=True, type='DefaultSampler'))
+train_pipeline = [
+    dict(backend_args=None, to_float32=True, type='LoadImageFromFile'),
+    dict(
+        backend_args=None,
+        type='LoadPanopticAnnotations',
+        with_bbox=True,
+        with_mask=True,
+        with_seg=True),
+    dict(prob=0.5, type='RandomFlip'),
+    dict(
+        keep_ratio=True,
+        ratio_range=(
+            0.1,
+            2.0,
+        ),
+        scale=(
+            1024,
+            1024,
+        ),
+        type='RandomResize'),
+    dict(
+        allow_negative_crop=True,
+        crop_size=(
+            1024,
+            1024,
+        ),
+        crop_type='absolute',
+        recompute_bbox=True,
+        type='RandomCrop'),
+    dict(type='PackDetInputs'),
+]
+val_cfg = dict(type='ValLoop')
+val_dataloader = dict(
+    batch_size=1,
+    dataset=dict(
+        ann_file='annotations/panoptic_val2017.json',
+        backend_args=None,
+        data_prefix=dict(img='val2017/', seg='annotations/panoptic_val2017/'),
+        data_root='data/coco/',
+        pipeline=[
+            dict(backend_args=None, type='LoadImageFromFile'),
+            dict(keep_ratio=True, scale=(
+                1333,
+                800,
+            ), type='Resize'),
+            dict(backend_args=None, type='LoadPanopticAnnotations'),
+            dict(
+                meta_keys=(
+                    'img_id',
+                    'img_path',
+                    'ori_shape',
+                    'img_shape',
+                    'scale_factor',
+                ),
+                type='PackDetInputs'),
+        ],
+        test_mode=True,
+        type='CocoPanopticDataset'),
+    drop_last=False,
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(shuffle=False, type='DefaultSampler'))
+val_evaluator = [
+    dict(
+        ann_file='data/coco/annotations/panoptic_val2017.json',
+        backend_args=None,
+        seg_prefix='data/coco/annotations/panoptic_val2017/',
+        type='CocoPanopticMetric'),
+    dict(
+        ann_file='data/coco/annotations/instances_val2017.json',
+        backend_args=None,
+        metric=[
+            'bbox',
+            'segm',
+        ],
+        type='CocoMetric'),
+]
+vis_backends = [
+    dict(type='LocalVisBackend'),
+]
+visualizer = dict(
+    name='visualizer',
+    type='DetLocalVisualizer',
+    vis_backends=[
+        dict(type='LocalVisBackend'),
+    ])

checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py

@@ -0,0 +1,350 @@
+auto_scale_lr = dict(base_batch_size=192)
+backend_args = dict(backend='local')
+codec = dict(
+    decode_center_shift=0.5,
+    decode_keypoint_order=[
+        0,
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        11,
+        12,
+        7,
+        8,
+        9,
+        10,
+        13,
+        14,
+        15,
+        16,
+    ],
+    decode_max_instances=30,
+    decode_topk=30,
+    heatmap_size=(
+        128,
+        128,
+    ),
+    input_size=(
+        512,
+        512,
+    ),
+    sigma=2,
+    type='AssociativeEmbedding')
+custom_hooks = [
+    dict(type='SyncBuffersHook'),
+]
+data_mode = 'bottomup'
+data_root = 'data/coco/'
+dataset_type = 'CocoDataset'
+default_hooks = dict(
+    badcase=dict(
+        badcase_thr=5,
+        enable=False,
+        metric_type='loss',
+        out_dir='badcase',
+        type='BadCaseAnalysisHook'),
+    checkpoint=dict(
+        interval=50,
+        rule='greater',
+        save_best='coco/AP',
+        type='CheckpointHook'),
+    logger=dict(interval=50, type='LoggerHook'),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+    timer=dict(type='IterTimerHook'),
+    visualization=dict(enable=False, type='PoseVisualizationHook'))
+default_scope = 'mmpose'
+env_cfg = dict(
+    cudnn_benchmark=False,
+    dist_cfg=dict(backend='nccl'),
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0))
+load_from = None
+log_level = 'INFO'
+log_processor = dict(
+    by_epoch=True, num_digits=6, type='LogProcessor', window_size=50)
+model = dict(
+    backbone=dict(
+        extra=dict(
+            stage1=dict(
+                block='BOTTLENECK',
+                num_blocks=(4, ),
+                num_branches=1,
+                num_channels=(64, ),
+                num_modules=1),
+            stage2=dict(
+                block='BASIC',
+                num_blocks=(
+                    4,
+                    4,
+                ),
+                num_branches=2,
+                num_channels=(
+                    32,
+                    64,
+                ),
+                num_modules=1),
+            stage3=dict(
+                block='BASIC',
+                num_blocks=(
+                    4,
+                    4,
+                    4,
+                ),
+                num_branches=3,
+                num_channels=(
+                    32,
+                    64,
+                    128,
+                ),
+                num_modules=4),
+            stage4=dict(
+                block='BASIC',
+                num_blocks=(
+                    4,
+                    4,
+                    4,
+                    4,
+                ),
+                num_branches=4,
+                num_channels=(
+                    32,
+                    64,
+                    128,
+                    256,
+                ),
+                num_modules=3)),
+        in_channels=3,
+        init_cfg=dict(
+            checkpoint=
+            'https://download.openmmlab.com/mmpose/pretrain_models/hrnet_w32-36af842e.pth',
+            type='Pretrained'),
+        type='HRNet'),
+    data_preprocessor=dict(
+        bgr_to_rgb=True,
+        mean=[
+            123.675,
+            116.28,
+            103.53,
+        ],
+        std=[
+            58.395,
+            57.12,
+            57.375,
+        ],
+        type='PoseDataPreprocessor'),
+    head=dict(
+        decoder=dict(
+            decode_center_shift=0.5,
+            decode_keypoint_order=[
+                0,
+                1,
+                2,
+                3,
+                4,
+                5,
+                6,
+                11,
+                12,
+                7,
+                8,
+                9,
+                10,
+                13,
+                14,
+                15,
+                16,
+            ],
+            decode_max_instances=30,
+            decode_topk=30,
+            heatmap_size=(
+                512,
+                512,
+            ),
+            input_size=(
+                512,
+                512,
+            ),
+            sigma=2,
+            type='AssociativeEmbedding'),
+        deconv_out_channels=None,
+        in_channels=32,
+        keypoint_loss=dict(type='KeypointMSELoss', use_target_weight=True),
+        num_keypoints=17,
+        tag_dim=1,
+        tag_loss=dict(loss_weight=0.001, type='AssociativeEmbeddingLoss'),
+        tag_per_keypoint=True,
+        type='AssociativeEmbeddingHead'),
+    test_cfg=dict(
+        align_corners=False,
+        flip_test=True,
+        multiscale_test=False,
+        restore_heatmap_size=True,
+        shift_heatmap=False),
+    type='BottomupPoseEstimator')
+optim_wrapper = dict(optimizer=dict(lr=0.0015, type='Adam'))
+param_scheduler = [
+    dict(
+        begin=0, by_epoch=False, end=500, start_factor=0.001, type='LinearLR'),
+    dict(
+        begin=0,
+        by_epoch=True,
+        end=300,
+        gamma=0.1,
+        milestones=[
+            200,
+            260,
+        ],
+        type='MultiStepLR'),
+]
+resume = False
+test_cfg = dict()
+test_dataloader = dict(
+    batch_size=1,
+    dataset=dict(
+        ann_file='annotations/person_keypoints_val2017.json',
+        data_mode='bottomup',
+        data_prefix=dict(img='val2017/'),
+        data_root='data/coco/',
+        pipeline=[
+            dict(type='LoadImage'),
+            dict(
+                input_size=(
+                    512,
+                    512,
+                ),
+                resize_mode='expand',
+                size_factor=64,
+                type='BottomupResize'),
+            dict(
+                meta_keys=(
+                    'id',
+                    'img_id',
+                    'img_path',
+                    'crowd_index',
+                    'ori_shape',
+                    'img_shape',
+                    'input_size',
+                    'input_center',
+                    'input_scale',
+                    'flip',
+                    'flip_direction',
+                    'flip_indices',
+                    'raw_ann_info',
+                    'skeleton_links',
+                ),
+                type='PackPoseInputs'),
+        ],
+        test_mode=True,
+        type='CocoDataset'),
+    drop_last=False,
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(round_up=False, shuffle=False, type='DefaultSampler'))
+test_evaluator = dict(
+    ann_file='data/coco/annotations/person_keypoints_val2017.json',
+    nms_mode='none',
+    score_mode='bbox',
+    type='CocoMetric')
+train_cfg = dict(by_epoch=True, max_epochs=300, val_interval=10)
+train_dataloader = dict(
+    batch_size=24,
+    dataset=dict(
+        ann_file='annotations/person_keypoints_train2017.json',
+        data_mode='bottomup',
+        data_prefix=dict(img='train2017/'),
+        data_root='data/coco/',
+        pipeline=[],
+        type='CocoDataset'),
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(shuffle=True, type='DefaultSampler'))
+train_pipeline = []
+val_cfg = dict()
+val_dataloader = dict(
+    batch_size=1,
+    dataset=dict(
+        ann_file='annotations/person_keypoints_val2017.json',
+        data_mode='bottomup',
+        data_prefix=dict(img='val2017/'),
+        data_root='data/coco/',
+        pipeline=[
+            dict(type='LoadImage'),
+            dict(
+                input_size=(
+                    512,
+                    512,
+                ),
+                resize_mode='expand',
+                size_factor=64,
+                type='BottomupResize'),
+            dict(
+                meta_keys=(
+                    'id',
+                    'img_id',
+                    'img_path',
+                    'crowd_index',
+                    'ori_shape',
+                    'img_shape',
+                    'input_size',
+                    'input_center',
+                    'input_scale',
+                    'flip',
+                    'flip_direction',
+                    'flip_indices',
+                    'raw_ann_info',
+                    'skeleton_links',
+                ),
+                type='PackPoseInputs'),
+        ],
+        test_mode=True,
+        type='CocoDataset'),
+    drop_last=False,
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(round_up=False, shuffle=False, type='DefaultSampler'))
+val_evaluator = dict(
+    ann_file='data/coco/annotations/person_keypoints_val2017.json',
+    nms_mode='none',
+    score_mode='bbox',
+    type='CocoMetric')
+val_pipeline = [
+    dict(type='LoadImage'),
+    dict(
+        input_size=(
+            512,
+            512,
+        ),
+        resize_mode='expand',
+        size_factor=64,
+        type='BottomupResize'),
+    dict(
+        meta_keys=(
+            'id',
+            'img_id',
+            'img_path',
+            'crowd_index',
+            'ori_shape',
+            'img_shape',
+            'input_size',
+            'input_center',
+            'input_scale',
+            'flip',
+            'flip_direction',
+            'flip_indices',
+            'raw_ann_info',
+            'skeleton_links',
+        ),
+        type='PackPoseInputs'),
+]
+vis_backends = [
+    dict(type='LocalVisBackend'),
+]
+visualizer = dict(
+    name='visualizer',
+    type='PoseLocalVisualizer',
+    vis_backends=[
+        dict(type='LocalVisBackend'),
+    ])

checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py

@@ -0,0 +1,336 @@
+backbone_norm_cfg = dict(requires_grad=True, type='LN')
+checkpoint_file = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/swin/swin_base_patch4_window7_224_22k_20220317-4f79f7c0.pth'
+crop_size = (
+    512,
+    512,
+)
+data_preprocessor = dict(
+    bgr_to_rgb=True,
+    mean=[
+        123.675,
+        116.28,
+        103.53,
+    ],
+    pad_val=0,
+    seg_pad_val=255,
+    size=(
+        512,
+        512,
+    ),
+    std=[
+        58.395,
+        57.12,
+        57.375,
+    ],
+    type='SegDataPreProcessor')
+data_root = 'data/ade/ADEChallengeData2016'
+dataset_type = 'ADE20KDataset'
+default_hooks = dict(
+    checkpoint=dict(by_epoch=False, interval=16000, type='CheckpointHook'),
+    logger=dict(interval=50, log_metric_by_epoch=False, type='LoggerHook'),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+    timer=dict(type='IterTimerHook'),
+    visualization=dict(type='SegVisualizationHook'))
+default_scope = 'mmseg'
+env_cfg = dict(
+    cudnn_benchmark=True,
+    dist_cfg=dict(backend='nccl'),
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0))
+img_ratios = [
+    0.5,
+    0.75,
+    1.0,
+    1.25,
+    1.5,
+    1.75,
+]
+load_from = None
+log_level = 'INFO'
+log_processor = dict(by_epoch=False)
+model = dict(
+    auxiliary_head=dict(
+        align_corners=False,
+        channels=256,
+        concat_input=False,
+        dropout_ratio=0.1,
+        in_channels=512,
+        in_index=2,
+        loss_decode=dict(
+            loss_weight=0.4, type='CrossEntropyLoss', use_sigmoid=False),
+        norm_cfg=dict(requires_grad=True, type='SyncBN'),
+        num_classes=150,
+        num_convs=1,
+        type='FCNHead'),
+    backbone=dict(
+        act_cfg=dict(type='GELU'),
+        attn_drop_rate=0.0,
+        depths=[
+            2,
+            2,
+            18,
+            2,
+        ],
+        drop_path_rate=0.3,
+        drop_rate=0.0,
+        embed_dims=128,
+        init_cfg=dict(
+            checkpoint=
+            'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/swin/swin_base_patch4_window7_224_22k_20220317-4f79f7c0.pth',
+            type='Pretrained'),
+        mlp_ratio=4,
+        norm_cfg=dict(requires_grad=True, type='LN'),
+        num_heads=[
+            4,
+            8,
+            16,
+            32,
+        ],
+        out_indices=(
+            0,
+            1,
+            2,
+            3,
+        ),
+        patch_norm=True,
+        patch_size=4,
+        pretrain_img_size=224,
+        qk_scale=None,
+        qkv_bias=True,
+        strides=(
+            4,
+            2,
+            2,
+            2,
+        ),
+        type='SwinTransformer',
+        use_abs_pos_embed=False,
+        window_size=7),
+    data_preprocessor=dict(
+        bgr_to_rgb=True,
+        mean=[
+            123.675,
+            116.28,
+            103.53,
+        ],
+        pad_val=0,
+        seg_pad_val=255,
+        size=(
+            512,
+            512,
+        ),
+        std=[
+            58.395,
+            57.12,
+            57.375,
+        ],
+        type='SegDataPreProcessor'),
+    decode_head=dict(
+        align_corners=False,
+        channels=512,
+        dropout_ratio=0.1,
+        in_channels=[
+            128,
+            256,
+            512,
+            1024,
+        ],
+        in_index=[
+            0,
+            1,
+            2,
+            3,
+        ],
+        loss_decode=dict(
+            loss_weight=1.0, type='CrossEntropyLoss', use_sigmoid=False),
+        norm_cfg=dict(requires_grad=True, type='SyncBN'),
+        num_classes=150,
+        pool_scales=(
+            1,
+            2,
+            3,
+            6,
+        ),
+        type='UPerHead'),
+    pretrained=None,
+    test_cfg=dict(mode='whole'),
+    train_cfg=dict(),
+    type='EncoderDecoder')
+norm_cfg = dict(requires_grad=True, type='SyncBN')
+optim_wrapper = dict(
+    optimizer=dict(
+        betas=(
+            0.9,
+            0.999,
+        ), lr=6e-05, type='AdamW', weight_decay=0.01),
+    paramwise_cfg=dict(
+        custom_keys=dict(
+            absolute_pos_embed=dict(decay_mult=0.0),
+            norm=dict(decay_mult=0.0),
+            relative_position_bias_table=dict(decay_mult=0.0))),
+    type='OptimWrapper')
+optimizer = dict(lr=0.01, momentum=0.9, type='SGD', weight_decay=0.0005)
+param_scheduler = [
+    dict(
+        begin=0, by_epoch=False, end=1500, start_factor=1e-06,
+        type='LinearLR'),
+    dict(
+        begin=1500,
+        by_epoch=False,
+        end=160000,
+        eta_min=0.0,
+        power=1.0,
+        type='PolyLR'),
+]
+resume = False
+test_cfg = dict(type='TestLoop')
+test_dataloader = dict(
+    batch_size=1,
+    dataset=dict(
+        data_prefix=dict(
+            img_path='images/validation',
+            seg_map_path='annotations/validation'),
+        data_root='data/ade/ADEChallengeData2016',
+        pipeline=[
+            dict(type='LoadImageFromFile'),
+            dict(keep_ratio=True, scale=(
+                2048,
+                512,
+            ), type='Resize'),
+            dict(reduce_zero_label=True, type='LoadAnnotations'),
+            dict(type='PackSegInputs'),
+        ],
+        type='ADE20KDataset'),
+    num_workers=4,
+    persistent_workers=True,
+    sampler=dict(shuffle=False, type='DefaultSampler'))
+test_evaluator = dict(
+    iou_metrics=[
+        'mIoU',
+    ], type='IoUMetric')
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(keep_ratio=True, scale=(
+        2048,
+        512,
+    ), type='Resize'),
+    dict(reduce_zero_label=True, type='LoadAnnotations'),
+    dict(type='PackSegInputs'),
+]
+train_cfg = dict(
+    max_iters=160000, type='IterBasedTrainLoop', val_interval=16000)
+train_dataloader = dict(
+    batch_size=2,
+    dataset=dict(
+        data_prefix=dict(
+            img_path='images/training', seg_map_path='annotations/training'),
+        data_root='data/ade/ADEChallengeData2016',
+        pipeline=[
+            dict(type='LoadImageFromFile'),
+            dict(reduce_zero_label=True, type='LoadAnnotations'),
+            dict(
+                keep_ratio=True,
+                ratio_range=(
+                    0.5,
+                    2.0,
+                ),
+                scale=(
+                    2048,
+                    512,
+                ),
+                type='RandomResize'),
+            dict(
+                cat_max_ratio=0.75, crop_size=(
+                    512,
+                    512,
+                ), type='RandomCrop'),
+            dict(prob=0.5, type='RandomFlip'),
+            dict(type='PhotoMetricDistortion'),
+            dict(type='PackSegInputs'),
+        ],
+        type='ADE20KDataset'),
+    num_workers=4,
+    persistent_workers=True,
+    sampler=dict(shuffle=True, type='InfiniteSampler'))
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(reduce_zero_label=True, type='LoadAnnotations'),
+    dict(
+        keep_ratio=True,
+        ratio_range=(
+            0.5,
+            2.0,
+        ),
+        scale=(
+            2048,
+            512,
+        ),
+        type='RandomResize'),
+    dict(cat_max_ratio=0.75, crop_size=(
+        512,
+        512,
+    ), type='RandomCrop'),
+    dict(prob=0.5, type='RandomFlip'),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='PackSegInputs'),
+]
+tta_model = dict(type='SegTTAModel')
+tta_pipeline = [
+    dict(backend_args=None, type='LoadImageFromFile'),
+    dict(
+        transforms=[
+            [
+                dict(keep_ratio=True, scale_factor=0.5, type='Resize'),
+                dict(keep_ratio=True, scale_factor=0.75, type='Resize'),
+                dict(keep_ratio=True, scale_factor=1.0, type='Resize'),
+                dict(keep_ratio=True, scale_factor=1.25, type='Resize'),
+                dict(keep_ratio=True, scale_factor=1.5, type='Resize'),
+                dict(keep_ratio=True, scale_factor=1.75, type='Resize'),
+            ],
+            [
+                dict(direction='horizontal', prob=0.0, type='RandomFlip'),
+                dict(direction='horizontal', prob=1.0, type='RandomFlip'),
+            ],
+            [
+                dict(type='LoadAnnotations'),
+            ],
+            [
+                dict(type='PackSegInputs'),
+            ],
+        ],
+        type='TestTimeAug'),
+]
+val_cfg = dict(type='ValLoop')
+val_dataloader = dict(
+    batch_size=1,
+    dataset=dict(
+        data_prefix=dict(
+            img_path='images/validation',
+            seg_map_path='annotations/validation'),
+        data_root='data/ade/ADEChallengeData2016',
+        pipeline=[
+            dict(type='LoadImageFromFile'),
+            dict(keep_ratio=True, scale=(
+                2048,
+                512,
+            ), type='Resize'),
+            dict(reduce_zero_label=True, type='LoadAnnotations'),
+            dict(type='PackSegInputs'),
+        ],
+        type='ADE20KDataset'),
+    num_workers=4,
+    persistent_workers=True,
+    sampler=dict(shuffle=False, type='DefaultSampler'))
+val_evaluator = dict(
+    iou_metrics=[
+        'mIoU',
+    ], type='IoUMetric')
+vis_backends = [
+    dict(type='LocalVisBackend'),
+]
+visualizer = dict(
+    name='visualizer',
+    type='SegLocalVisualizer',
+    vis_backends=[
+        dict(type='LocalVisBackend'),
+    ])

configs/ctaic/s1_det_instance.yaml

@@ -2,7 +2,7 @@
 root: "logs"
 exp_name: "ctaic_s1"
 scenario: "s1"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 # trained base-layer TAIC for detection (provides y_b_hat)
 base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar"
@@ -24,3 +24,8 @@ seed: 42
 pretrained_backbone: true
 out_channels: 128  # extension = Cascade Mask R-CNN + Swin-B
 align_mode: "fpn"
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth"

configs/ctaic/s2_sem_panoptic.yaml

@@ -2,7 +2,7 @@
 root: "logs"
 exp_name: "ctaic_s2"
 scenario: "s2"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 base_taic_checkpoint: "./checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar"
 taic_init: "./checkpoints/taic/panoptic/1/checkpoint_best_loss.pth.tar"
@@ -20,3 +20,8 @@ cuda: true
 save: true
 seed: 42
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py"
+task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth"

configs/ctaic/s3_det_pose.yaml

@@ -2,7 +2,7 @@
 root: "logs"
 exp_name: "ctaic_s3"
 scenario: "s3"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar"
 taic_init: "./checkpoints/taic/pose/1/checkpoint_best_loss.pth.tar"
@@ -22,3 +22,8 @@ seed: 42
 pretrained_backbone: true
 out_channels: 32
 align_mode: "stages"
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py"
+task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth"

configs/eval/ctaic_s1.yaml

@@ -1,5 +1,5 @@
 scenario: "s1"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/instances_val2017.json"
 quality_level: 1
@@ -7,10 +7,13 @@ lmbda: 0.0035
 checkpoint: "./checkpoints/ctaic/s1_det_instance/stage2/1/checkpoint_best_loss.pth.tar"
 base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
-task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
-task_checkpoint: "./checkpoints/task_networks/instance/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth"

configs/eval/ctaic_s2.yaml

@@ -1,5 +1,5 @@
 scenario: "s2"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/panoptic_val2017.json"
 quality_level: 1
@@ -7,11 +7,14 @@ lmbda: 0.0035
 checkpoint: "./checkpoints/ctaic/s2_sem_panoptic/stage2/1/checkpoint_best_loss.pth.tar"
 base_taic_checkpoint: "./checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
-task_config: "./configs/task_networks/maskformer_swin-b_coco.py"
-task_checkpoint: "./checkpoints/task_networks/panoptic/model.pth"
 panoptic_gt_folder: "/data/coco2017/annotations/panoptic_val2017"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py"
+task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth"

configs/eval/ctaic_s3.yaml

@@ -1,5 +1,5 @@
 scenario: "s3"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/coco_wholebody_val_v1.0.json"
 quality_level: 1
@@ -7,10 +7,13 @@ lmbda: 0.0035
 checkpoint: "./checkpoints/ctaic/s3_det_pose/stage2/1/checkpoint_best_loss.pth.tar"
 base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
-task_config: "./configs/task_networks/higherhrnet_w32_coco_wholebody.py"
-task_checkpoint: "./checkpoints/task_networks/pose/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py"
+task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth"

configs/eval/taic_detection.yaml

@@ -1,17 +1,20 @@
 # TAIC eval: codec test + optional --with-metrics
 task: "detection"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/instances_val2017.json"
 quality_level: 1
 lmbda: 0.0035
-checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar"
+checkpoint: "./logs/taic_detection/1/checkpoint_0.1040.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 # official task network (required for --with-metrics)
-task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
-task_checkpoint: "./checkpoints/task_networks/detection/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/detection/model_mmdet3.pth"

configs/eval/taic_instance.yaml

@@ -1,6 +1,6 @@
 # TAIC eval: codec test + optional --with-metrics
 task: "instance"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/instances_val2017.json"
 quality_level: 1
@@ -8,10 +8,13 @@ lmbda: 0.0035
 checkpoint: "./checkpoints/taic/instance/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 # official task network (required for --with-metrics)
-task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
-task_checkpoint: "./checkpoints/task_networks/instance/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth"

configs/eval/taic_panoptic.yaml

@@ -1,6 +1,6 @@
 # TAIC eval: codec test + optional --with-metrics
 task: "panoptic"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/panoptic_val2017.json"
 quality_level: 1
@@ -8,10 +8,13 @@ lmbda: 0.0035
 checkpoint: "./checkpoints/taic/panoptic/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 # official task network (required for --with-metrics)
-task_config: "./configs/task_networks/maskformer_swin-b_coco.py"
-task_checkpoint: "./checkpoints/task_networks/panoptic/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py"
+task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth"

configs/eval/taic_pose.yaml

@@ -1,6 +1,6 @@
 # TAIC eval: codec test + optional --with-metrics
 task: "pose"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/coco_wholebody_val_v1.0.json"
 quality_level: 1
@@ -8,10 +8,13 @@ lmbda: 0.0035
 checkpoint: "./checkpoints/taic/pose/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 # official task network (required for --with-metrics)
-task_config: "./configs/task_networks/higherhrnet_w32_coco_wholebody.py"
-task_checkpoint: "./checkpoints/task_networks/pose/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py"
+task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth"

configs/eval/taic_semantic.yaml

@@ -1,6 +1,6 @@
 # TAIC eval: codec test + optional --with-metrics
 task: "semantic"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 split: "val2017"
 ann_file: "annotations/panoptic_val2017.json"
 quality_level: 1
@@ -8,10 +8,13 @@ lmbda: 0.0035
 checkpoint: "./checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 # official task network (required for --with-metrics)
-task_config: "./configs/task_networks/upernet_swin-b_coco.py"
-task_checkpoint: "./checkpoints/task_networks/semantic/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1
 num_workers: 4
 pretrained_backbone: true
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py"
+task_checkpoint: "./checkpoints/task_networks/semantic/upernet_swin_base_patch4_window7_512x512_160k_ade20k_pretrain_224x224_22K_20210526_211650-762e2178.pth"

configs/taic/detection.yaml

@@ -2,17 +2,18 @@
 root: "logs"
 exp_name: "taic_detection"
 task: "detection"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"  # quality matches lmbda mapping
-checkpoint: null
-epochs: 35
+checkpoint: null # "./logs/taic_detection/1/checkpoint_60.pth.tar"
+freeze_base_codec: false  # false = fine-tune the entire TIC + SFMA + Task Connector
+epochs: 999
 learning_rate: 1.0e-4
 gpu_id: 0
 quality_level: 1
 # paper lambda set: {0.0035, 0.0067, 0.0130, 0.0250}
-lmbda: 0.0035
+lmbda: 0.2050
 num_workers: 8
-batch_size: 80
+batch_size: 32
 test_batch_size: 1
 patch_size: 256
 cuda: true
@@ -21,3 +22,8 @@ seed: 42
 pretrained_backbone: true
 out_channels: 128
 align_mode: "fpn"
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/detection/model_mmdet3.pth"

configs/taic/instance.yaml

@@ -2,7 +2,7 @@
 root: "logs"
 exp_name: "taic_instance"
 task: "instance"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 checkpoint: null
 epochs: 35
@@ -20,3 +20,8 @@ seed: 42
 pretrained_backbone: true
 out_channels: 128  # Cascade Mask R-CNN + Swin-B F1
 align_mode: "fpn"
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth"

configs/taic/panoptic.yaml

@@ -1,7 +1,7 @@
 root: "logs"
 exp_name: "taic_panoptic"
 task: "panoptic"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 checkpoint: null
 epochs: 35
@@ -19,3 +19,8 @@ seed: 42
 pretrained_backbone: true
 out_channels: 128
 align_mode: "stages"
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py"
+task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth"

configs/taic/pose.yaml

@@ -1,7 +1,7 @@
 root: "logs"
 exp_name: "taic_pose"
 task: "pose"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 checkpoint: null
 epochs: 35
@@ -19,3 +19,8 @@ seed: 42
 pretrained_backbone: true
 out_channels: 32
 align_mode: "stages"
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py"
+task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth"

configs/taic/semantic.yaml

@@ -1,7 +1,7 @@
 root: "logs"
 exp_name: "taic_semantic"
 task: "semantic"
-dataset_path: "/data/coco2017"
+dataset_path: "/data/Dataset/coco2017"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
 checkpoint: null
 epochs: 35
@@ -19,3 +19,8 @@ seed: 42
 pretrained_backbone: true
 out_channels: 128
 align_mode: "fpn"
+
+# Official task-network teacher (backbone/neck) for feature alignment D
+use_official_teacher: true
+task_config: "./checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py"
+task_checkpoint: "./checkpoints/task_networks/semantic/upernet_swin_base_patch4_window7_512x512_160k_ade20k_pretrain_224x224_22K_20210526_211650-762e2178.pth"

configs/task_networks/cascade_mask_rcnn_swin_base_coco.py

@@ -1,19 +1,251 @@
-# STUB / pointer config for Cascade Mask R-CNN + Swin-B (official Swin detection zoo).
-#
-# Source repo: https://github.com/SwinTransformer/Swin-Transformer-Object-Detection
-# Official config:
-#   configs/swin/cascade_mask_rcnn_swin_base_patch4_window7_mstrain_480-800_giou_4conv1f_adamw_3x_coco.py
-# Official checkpoint:
-#   https://github.com/SwinTransformer/storage/releases/download/v1.0.2/cascade_mask_rcnn_swin_base_patch4_window7.pth
-#
-# Replace this file by copying/symlinking the real config from that repo, then set
-# task_checkpoint to the downloaded .pth under checkpoints/task_networks/detection/.
-#
-# Used for FlexICM detection (mAP-bbox) and instance segmentation (mAP-mask).
-# Backbone F1 channels = 128 (Swin-B).
+# Cascade Mask R-CNN + Swin-B for FlexICM metric eval (mmdet 3.x).
+# Matches the official Swin-Transformer-Object-Detection paper checkpoint:
+#   cascade_mask_rcnn_swin_base_patch4_window7.pth
+# Use the converted weights:
+#   checkpoints/task_networks/detection/model_mmdet3.pth
+# (produced by scripts/convert_swin_det_ckpt_to_mmdet3.py)
 
-raise RuntimeError(
-    "Replace configs/task_networks/cascade_mask_rcnn_swin_base_coco.py with the official "
-    "Swin-Transformer-Object-Detection config: "
-    "cascade_mask_rcnn_swin_base_patch4_window7_mstrain_480-800_giou_4conv1f_adamw_3x_coco.py"
-)
+_base_ = []  # self-contained for FlexICM
+
+model = dict(
+    type='CascadeRCNN',
+    data_preprocessor=dict(
+        type='DetDataPreprocessor',
+        mean=[123.675, 116.28, 103.53],
+        std=[58.395, 57.12, 57.375],
+        bgr_to_rgb=True,
+        pad_mask=True,
+        pad_size_divisor=32),
+    backbone=dict(
+        type='SwinTransformer',
+        embed_dims=128,
+        depths=[2, 2, 18, 2],
+        num_heads=[4, 8, 16, 32],
+        window_size=7,
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.3,
+        patch_norm=True,
+        out_indices=(0, 1, 2, 3),
+        with_cp=False,
+        convert_weights=False,
+        init_cfg=None),
+    neck=dict(
+        type='FPN',
+        in_channels=[128, 256, 512, 1024],
+        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=3,
+        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(
+                type='Shared4Conv1FCBBoxHead',
+                in_channels=256,
+                conv_out_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0.0, 0.0, 0.0, 0.0],
+                    target_stds=[0.1, 0.1, 0.2, 0.2]),
+                reg_class_agnostic=False,
+                reg_decoded_bbox=True,
+                norm_cfg=dict(type='BN', requires_grad=True),
+                loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+                loss_bbox=dict(type='GIoULoss', loss_weight=10.0)),
+            dict(
+                type='Shared4Conv1FCBBoxHead',
+                in_channels=256,
+                conv_out_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0.0, 0.0, 0.0, 0.0],
+                    target_stds=[0.05, 0.05, 0.1, 0.1]),
+                reg_class_agnostic=False,
+                reg_decoded_bbox=True,
+                norm_cfg=dict(type='BN', requires_grad=True),
+                loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+                loss_bbox=dict(type='GIoULoss', loss_weight=10.0)),
+            dict(
+                type='Shared4Conv1FCBBoxHead',
+                in_channels=256,
+                conv_out_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0.0, 0.0, 0.0, 0.0],
+                    target_stds=[0.033, 0.033, 0.067, 0.067]),
+                reg_class_agnostic=False,
+                reg_decoded_bbox=True,
+                norm_cfg=dict(type='BN', requires_grad=True),
+                loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+                loss_bbox=dict(type='GIoULoss', loss_weight=10.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)),
+            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)),
+            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)),
+        ]),
+    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),
+                mask_size=28,
+                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),
+                mask_size=28,
+                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),
+                mask_size=28,
+                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,
+            mask_thr_binary=0.5)))
+
+# Minimal dataset meta required by mmdet.apis.init_detector
+test_dataloader = dict(
+    dataset=dict(
+        type='CocoDataset',
+        ann_file='annotations/instances_val2017.json',
+        data_prefix=dict(img='val2017/'),
+        metainfo=dict(
+            classes=('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
+                     'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
+                     'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
+                     'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
+                     'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
+                     'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
+                     'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
+                     'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
+                     'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
+                     'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
+                     'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
+                     'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
+                     'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
+                     'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'))))

flexicm/models/taic.py

@@ -374,7 +374,7 @@ class TAIC(nn.Module):
             ["_quantized_cdf", "_offset", "_cdf_length", "scale_table"],
             state_dict,
         )
-        super().load_state_dict(state_dict, strict=strict)
+        return super().load_state_dict(state_dict, strict=strict)
 
     def load_base_codec(self, state_dict, strict=False):
         """Load pretrained TIC / TIC-SFMA weights into matching modules."""

flexicm/tasks/__init__.py

@@ -21,6 +21,41 @@ import torch.nn.functional as F
 
 from flexicm.tasks.losses import freeze_module
 from flexicm.tasks.swin_teacher import SwinStageTeacher
+from flexicm.tasks.official_teachers import OfficialHRNetTeacher, OfficialSwinTeacher
+
+# Default official task-network assets (used when use_official_teacher=True)
+DEFAULT_TEACHER_ASSETS = {
+    "detection": {
+        "task_config": "configs/task_networks/cascade_mask_rcnn_swin_base_coco.py",
+        "task_checkpoint": "checkpoints/task_networks/detection/model_mmdet3.pth",
+        "framework": "mmdet",
+        "align_mode": "fpn",
+    },
+    "instance": {
+        "task_config": "configs/task_networks/cascade_mask_rcnn_swin_base_coco.py",
+        "task_checkpoint": "checkpoints/task_networks/instance/model_mmdet3.pth",
+        "framework": "mmdet",
+        "align_mode": "fpn",
+    },
+    "semantic": {
+        "task_config": "checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py",
+        "task_checkpoint": "checkpoints/task_networks/semantic/upernet_swin_base_patch4_window7_512x512_160k_ade20k_pretrain_224x224_22K_20210526_211650-762e2178.pth",
+        "framework": "mmseg",
+        "align_mode": "fpn",
+    },
+    "panoptic": {
+        "task_config": "checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py",
+        "task_checkpoint": "checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth",
+        "framework": "mmdet",
+        "align_mode": "stages",
+    },
+    "pose": {
+        "task_config": "checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py",
+        "task_checkpoint": "checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth",
+        "framework": "mmpose",
+        "align_mode": "stages",
+    },
+}
 
 
 class DetectionTeacher(nn.Module):
@@ -179,15 +214,64 @@ class HigherHRNetTeacher(nn.Module):
 def build_teacher(task: str, **kwargs) -> nn.Module:
     task = task.lower()
     pretrained = kwargs.pop("pretrained_backbone", kwargs.pop("pretrained", True))
-    if task in ("detection", "object_detection", "det"):
+    use_official = kwargs.pop("use_official_teacher", True)
+    task_config = kwargs.pop("task_config", None)
+    task_checkpoint = kwargs.pop("task_checkpoint", None)
+    device = kwargs.pop("device", "cpu")
+
+    # Normalize aliases
+    alias = {
+        "object_detection": "detection",
+        "det": "detection",
+        "instance_seg": "instance",
+        "instance_segmentation": "instance",
+        "semantic_seg": "semantic",
+        "semantic_segmentation": "semantic",
+        "panoptic_seg": "panoptic",
+        "panoptic_segmentation": "panoptic",
+        "pose_estimation": "pose",
+    }
+    task = alias.get(task, task)
+
+    if use_official:
+        assets = DEFAULT_TEACHER_ASSETS.get(task)
+        if assets is None:
+            raise ValueError(f"Unknown task for official teacher: {task}")
+        cfg = task_config or assets["task_config"]
+        ckpt = task_checkpoint or assets["task_checkpoint"]
+        align_mode = kwargs.pop("align_mode", assets["align_mode"])
+        framework = assets["framework"]
+        if framework in ("mmdet", "mmseg"):
+            teacher = OfficialSwinTeacher(
+                config_path=cfg,
+                checkpoint_path=ckpt,
+                align_mode=align_mode,
+                framework=framework,
+                device=device,
+            )
+            freeze_module(teacher)
+            return teacher
+        if framework == "mmpose":
+            width = kwargs.pop("width", 32)
+            teacher = OfficialHRNetTeacher(
+                config_path=cfg,
+                checkpoint_path=ckpt,
+                device=device,
+                width=width,
+            )
+            freeze_module(teacher)
+            return teacher
+
+    # Legacy timm / in-repo teachers (ImageNet Swin / stem HRNet)
+    if task == "detection":
         return DetectionTeacher(task="detection", pretrained_backbone=pretrained, **kwargs)
-    if task in ("instance", "instance_seg", "instance_segmentation"):
+    if task == "instance":
         return DetectionTeacher(task="instance", pretrained_backbone=pretrained, **kwargs)
-    if task in ("semantic", "semantic_seg", "semantic_segmentation"):
+    if task == "semantic":
         return SemanticSegTeacher(pretrained_backbone=pretrained, **kwargs)
-    if task in ("panoptic", "panoptic_seg", "panoptic_segmentation"):
+    if task == "panoptic":
         return PanopticSegTeacher(pretrained_backbone=pretrained, **kwargs)
-    if task in ("pose", "pose_estimation"):
+    if task == "pose":
         return HigherHRNetTeacher(pretrained=pretrained, **kwargs)
     raise ValueError(f"Unknown task: {task}")
 

flexicm/tasks/metric_eval.py

@@ -73,7 +73,10 @@ def run_task_metric_eval(
                 else:
                     out = codec(x)
 
-            h = crop_feature_to_image(out["h"], (H, W))
+            h = out["h"]
+            meta = dict(meta)
+            meta["pad_height"] = int(x.shape[-2])
+            meta["pad_width"] = int(x.shape[-1])
             pred = runner.predict_from_h(h, meta)
             predictions.append(pred)
 

flexicm/tasks/metric_runners.py

@@ -85,7 +85,9 @@ def _require_mmpose():
 def swin_feats_from_h(backbone: nn.Module, h: torch.Tensor) -> Tuple[torch.Tensor, ...]:
     """Treat h as Swin F1 (stage-0 output) and run remaining stages.
 
-    Compatible with MMDet/MMSeg SwinTransformer that exposes `.stages` / `.layers`.
+    Compatible with:
+      - MMDet 3.x SwinTransformer (token + hw_shape API, ``stages``)
+      - older Swin / MMSeg variants that expose ``stages`` / ``layers``
     """
     stages = None
     for name in ("stages", "layers"):
@@ -95,36 +97,77 @@ def swin_feats_from_h(backbone: nn.Module, h: torch.Tensor) -> Tuple[torch.Tenso
     if stages is None:
         raise RuntimeError("Backbone has no stages/layers; cannot inject h as F1")
 
+    # h is NCHW (B,C,H,W) == F1
+    if h.dim() != 4:
+        raise ValueError(f"expected NCHW h, got shape {tuple(h.shape)}")
+    B, C, H, W = h.shape
     outs = [h]
+
+    # Detect mmdet3-style SwinBlockSequence: forward(x, hw_shape)
+    import inspect
+
+    try:
+        needs_hw = "hw_shape" in inspect.signature(stages[0].forward).parameters
+    except (TypeError, ValueError):
+        needs_hw = False
+
+    if needs_hw:
+        # tokens: (B, H*W, C)
+        x = h.flatten(2).transpose(1, 2).contiguous()
+        hw_shape = (H, W)
+
+        # F1 is stage-0 block output (pre-downsample). Feed stage-0 downsample
+        # then run stages 1..N-1, mirroring SwinTransformer.forward.
+        if getattr(stages[0], "downsample", None) is not None:
+            x, hw_shape = stages[0].downsample(x, hw_shape)
+
+        for i in range(1, len(stages)):
+            x, hw_shape, out, out_hw_shape = stages[i](x, hw_shape)
+            norm_name = f"norm{i}"
+            if hasattr(backbone, norm_name):
+                out = getattr(backbone, norm_name)(out)
+            feat_dim = out.shape[-1]
+            out = (
+                out.view(B, out_hw_shape[0], out_hw_shape[1], feat_dim)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
+            outs.append(out)
+
+        # Optionally normalize F1 with norm0 for consistency with extract_feat
+        if hasattr(backbone, "norm0"):
+            f1 = h.flatten(2).transpose(1, 2).contiguous()
+            f1 = backbone.norm0(f1)
+            outs[0] = (
+                f1.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
+            )
+        return tuple(outs)
+
+    # Legacy path (timm / older mmdet): stage modules accept NCHW / plain tensors
     x = h
-    # Stage 0 already produced F1 (=h); run stages 1..N-1
     for i in range(1, len(stages)):
         x = stages[i](x)
         if isinstance(x, (tuple, list)):
             x = x[0]
-        # MMDet Swin may return NCHW already
-        if x.dim() == 4 and x.shape[1] < x.shape[-1] and x.shape[-1] in (96, 128, 192, 256, 384, 512, 768, 1024):
-            # likely NHWC
+        if x.dim() == 4 and x.shape[1] < x.shape[-1] and x.shape[-1] in (
+            96, 128, 192, 256, 384, 512, 768, 1024
+        ):
             x = x.permute(0, 3, 1, 2).contiguous()
         outs.append(x)
 
-    # Apply per-stage norms if present (mmdet Swin)
     if hasattr(backbone, "num_features") or hasattr(backbone, "out_indices"):
         norm_outs = []
         for i, out in enumerate(outs):
             norm_name = f"norm{i}"
             if hasattr(backbone, norm_name):
                 nchw = out
-                # LayerNorm over channel last in some impls
                 norm = getattr(backbone, norm_name)
-                if nchw.shape[1] == getattr(norm, "normalized_shape", [nchw.shape[1]])[0] if hasattr(norm, "normalized_shape") else True:
-                    # try NCHW LayerNorm via transpose
-                    try:
-                        y = nchw.permute(0, 2, 3, 1)
-                        y = norm(y)
-                        nchw = y.permute(0, 3, 1, 2).contiguous()
-                    except Exception:
-                        nchw = out
+                try:
+                    y = nchw.permute(0, 2, 3, 1)
+                    y = norm(y)
+                    nchw = y.permute(0, 3, 1, 2).contiguous()
+                except Exception:
+                    nchw = out
                 norm_outs.append(nchw)
             else:
                 norm_outs.append(out)
@@ -149,39 +192,39 @@ class DetectionMetricRunner(TaskMetricRunner):
     @torch.no_grad()
     def predict_from_h(self, h: torch.Tensor, img_meta: Dict[str, Any]) -> Dict[str, Any]:
         assert self.model is not None
-        # h: 1xCx(H/4)x(W/4)
+        # h: 1xCx(H/4)x(W/4) — preferably on the padded codec grid
         backbone = self.model.backbone
         feats = swin_feats_from_h(backbone, h)
         if hasattr(self.model, "neck") and self.model.neck is not None:
             feats = self.model.neck(feats)
 
         # Build a minimal img_metas / data_samples for mmdet 3.x or 2.x
-        H, W = int(img_meta["height"]), int(img_meta["width"])
+        ori_h, ori_w = int(img_meta["height"]), int(img_meta["width"])
+        # If h comes from a padded codec input, prefer those spatial sizes so
+        # FPN strides line up; boxes are rescaled back to ori_shape.
+        pad_h = int(img_meta.get("pad_height", h.shape[-2] * 4))
+        pad_w = int(img_meta.get("pad_width", h.shape[-1] * 4))
         try:
             # MMDet 3.x style
             from mmdet.structures import DetDataSample
-            from mmengine.structures import InstanceData
 
             data_sample = DetDataSample()
             data_sample.set_metainfo(
                 dict(
-                    img_shape=(H, W),
-                    ori_shape=(H, W),
-                    pad_shape=(H, W),
+                    img_shape=(pad_h, pad_w),
+                    ori_shape=(ori_h, ori_w),
+                    pad_shape=(pad_h, pad_w),
                     scale_factor=(1.0, 1.0),
                     img_id=img_meta.get("image_id"),
                 )
             )
-            # Use RPN + ROI heads with injected feats
-            if hasattr(self.model, "extract_feat"):
-                # bypass extract_feat by calling predict with feats if supported
-                pass
             rpn_results_list = self.model.rpn_head.predict(feats, [data_sample], rescale=False)
             results_list = self.model.roi_head.predict(
                 feats, rpn_results_list, [data_sample], rescale=True
             )
             pred = results_list[0]
-            inst = pred.pred_instances
+            # mmdet 3.x roi_head.predict may return DetDataSample or InstanceData
+            inst = pred.pred_instances if hasattr(pred, "pred_instances") else pred
             out = {
                 "image_id": img_meta["image_id"],
                 "bboxes": inst.bboxes.detach().cpu(),
@@ -195,9 +238,9 @@ class DetectionMetricRunner(TaskMetricRunner):
             # Fallback MMDet 2.x
             img_metas = [
                 dict(
-                    img_shape=(H, W, 3),
-                    ori_shape=(H, W, 3),
-                    pad_shape=(H, W, 3),
+                    img_shape=(pad_h, pad_w, 3),
+                    ori_shape=(ori_h, ori_w, 3),
+                    pad_shape=(pad_h, pad_w, 3),
                     scale_factor=1.0,
                     flip=False,
                 )

flexicm/tasks/official_teachers.py

@@ -0,0 +1,279 @@
+"""Teachers that load official task-network backbones (+ necks) for feature alignment.
+
+Training Distortion D is computed against these frozen features so that codec
+outputs match the same backbone used at metric evaluation time.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from flexicm.tasks.losses import freeze_module
+from flexicm.tasks.metric_runners import swin_feats_from_h
+
+
+def _resolve_path(path: Optional[str], repo_root: Optional[str] = None) -> Optional[str]:
+    if not path:
+        return None
+    if os.path.isabs(path):
+        return path
+    if repo_root is None:
+        repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+    return os.path.join(repo_root, path)
+
+
+def _imagenet_norm_rgb01(x: torch.Tensor) -> torch.Tensor:
+    """Normalize RGB float images in [0, 1] with ImageNet mean/std (float space)."""
+    mean = x.new_tensor([0.485, 0.456, 0.406])[None, :, None, None]
+    std = x.new_tensor([0.229, 0.224, 0.225])[None, :, None, None]
+    return (x - mean) / std
+
+
+def _mmdet_preprocess_rgb01(x: torch.Tensor, model: nn.Module) -> torch.Tensor:
+    """Apply DetDataPreprocessor-equivalent norm: RGB [0,1] -> model input."""
+    pp = getattr(model, "data_preprocessor", None)
+    if pp is None:
+        return _imagenet_norm_rgb01(x)
+    # DetDataPreprocessor stores mean/std for 0-255 inputs
+    mean = getattr(pp, "mean", None)
+    std = getattr(pp, "std", None)
+    if mean is None or std is None:
+        return _imagenet_norm_rgb01(x)
+    mean_t = mean.view(1, -1, 1, 1).to(device=x.device, dtype=x.dtype)
+    std_t = std.view(1, -1, 1, 1).to(device=x.device, dtype=x.dtype)
+    # x in [0,1] RGB; preprocessor mean/std are for 0-255 RGB after bgr_to_rgb
+    return (x * 255.0 - mean_t) / std_t
+
+
+class OfficialSwinTeacher(nn.Module):
+    """Frozen Swin teacher from an official MMDet / MMSeg checkpoint.
+
+    align_mode:
+      - ``fpn``: P2..P6 via backbone + FPN/neck (detection / instance / semantic)
+      - ``stages``: F1..F4 via backbone stages (panoptic)
+    """
+
+    def __init__(
+        self,
+        config_path: str,
+        checkpoint_path: str,
+        align_mode: str = "fpn",
+        framework: str = "mmdet",
+        device: str = "cpu",
+    ):
+        super().__init__()
+        assert align_mode in ("fpn", "stages")
+        self.align_mode = align_mode
+        self.framework = framework
+        self.config_path = _resolve_path(config_path)
+        self.checkpoint_path = _resolve_path(checkpoint_path)
+        if not self.config_path or not os.path.isfile(self.config_path):
+            raise FileNotFoundError(f"task_config not found: {config_path}")
+        if not self.checkpoint_path or not os.path.isfile(self.checkpoint_path):
+            raise FileNotFoundError(f"task_checkpoint not found: {checkpoint_path}")
+
+        self.model = self._load_model(device)
+        freeze_module(self.model)
+        self.out_channels = self._infer_f1_channels()
+        self._aux_fpn = None
+
+    def _load_model(self, device: str) -> nn.Module:
+        if self.framework == "mmdet":
+            from mmdet.apis import init_detector
+
+            return init_detector(self.config_path, self.checkpoint_path, device=device)
+        if self.framework == "mmseg":
+            from mmseg.apis import init_model
+
+            return init_model(self.config_path, self.checkpoint_path, device=device)
+        raise ValueError(f"Unknown framework={self.framework}")
+
+    def _infer_f1_channels(self) -> int:
+        bb = self.model.backbone
+        if hasattr(bb, "num_features"):
+            return int(bb.num_features[0])
+        if hasattr(bb, "embed_dims"):
+            return int(bb.embed_dims)
+        return 128
+
+    def _preprocess(self, images: torch.Tensor) -> torch.Tensor:
+        return _mmdet_preprocess_rgb01(images, self.model)
+
+    def _backbone_stages_from_image(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        """Return F1..F4 NCHW from Swin backbone (with per-stage norms)."""
+        bb = self.model.backbone
+        # Prefer official forward when possible (handles abs pos / dropout)
+        if hasattr(bb, "forward") and not hasattr(bb, "stages"):
+            feats = bb(x)
+            return tuple(feats[:4])
+
+        tokens, hw_shape = bb.patch_embed(x)
+        if getattr(bb, "use_abs_pos_embed", False):
+            tokens = tokens + bb.absolute_pos_embed
+        tokens = bb.drop_after_pos(tokens)
+        outs = []
+        for i, stage in enumerate(bb.stages):
+            tokens, hw_shape, out, out_hw = stage(tokens, hw_shape)
+            if i in getattr(bb, "out_indices", (0, 1, 2, 3)):
+                norm = getattr(bb, f"norm{i}")
+                out = norm(out)
+                feat = (
+                    out.view(-1, *out_hw, bb.num_features[i])
+                    .permute(0, 3, 1, 2)
+                    .contiguous()
+                )
+                outs.append(feat)
+        return tuple(outs)
+
+    def _fpn_from_stages(self, stages: Tuple[torch.Tensor, ...]) -> Dict[str, torch.Tensor]:
+        neck = getattr(self.model, "neck", None)
+        if neck is None and self.framework == "mmseg":
+            # UPerNet has no separate FPN neck; reuse Cascade Mask R-CNN FPN
+            # (same Swin-B channel layout) so Eq.2 aligns P2..P6.
+            neck = self._get_or_build_aux_fpn(stages[0].device)
+        if neck is None:
+            p = list(stages[:4])
+            while len(p) < 4:
+                p.append(p[-1])
+            p2, p3, p4, p5 = p
+            p6 = F.avg_pool2d(p5, kernel_size=2, stride=2)
+            return {"p2": p2, "p3": p3, "p4": p4, "p5": p5, "p6": p6}
+
+        pyramid = neck(stages)
+        keys = ["p2", "p3", "p4", "p5", "p6"]
+        out = {}
+        for i, feat in enumerate(pyramid):
+            if i < len(keys):
+                out[keys[i]] = feat
+        return out
+
+    def _get_or_build_aux_fpn(self, device) -> nn.Module:
+        if hasattr(self, "_aux_fpn") and self._aux_fpn is not None:
+            return self._aux_fpn
+        from mmdet.models.necks import FPN
+
+        fpn = FPN(
+            in_channels=[128, 256, 512, 1024],
+            out_channels=256,
+            num_outs=5,
+        )
+        # Load official Cascade Mask R-CNN neck weights when available
+        det_ckpt = _resolve_path("checkpoints/task_networks/detection/model_mmdet3.pth")
+        if det_ckpt and os.path.isfile(det_ckpt):
+            raw = torch.load(det_ckpt, map_location="cpu")
+            state = raw["state_dict"] if isinstance(raw, dict) and "state_dict" in raw else raw
+            neck_state = {
+                k[len("neck.") :]: v for k, v in state.items() if k.startswith("neck.")
+            }
+            missing, unexpected = fpn.load_state_dict(neck_state, strict=False)
+            # missing/unexpected are fine to ignore for BN tracking buffers etc.
+        self._aux_fpn = freeze_module(fpn.to(device))
+        return self._aux_fpn
+
+    def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]:
+        x = self._preprocess(images)
+        if (
+            self.align_mode == "fpn"
+            and hasattr(self.model, "extract_feat")
+            and getattr(self.model, "neck", None) is not None
+        ):
+            # CascadeRCNN.extract_feat = neck(backbone(x)) -> P2..P6
+            feats = self.model.extract_feat(x)
+            keys = ["p2", "p3", "p4", "p5", "p6"]
+            return {keys[i]: feats[i] for i in range(min(len(keys), len(feats)))}
+
+        if self.framework == "mmseg" and hasattr(self.model, "extract_feat"):
+            stages = tuple(self.model.extract_feat(x)[:4])
+        else:
+            stages = self._backbone_stages_from_image(x)
+        if self.align_mode == "stages":
+            return {f"f{i+1}": stages[i] for i in range(min(4, len(stages)))}
+        return self._fpn_from_stages(stages)
+
+    def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]:
+        stages = swin_feats_from_h(self.model.backbone, h)
+        if self.align_mode == "stages":
+            return {f"f{i+1}": stages[i] for i in range(min(4, len(stages)))}
+        return self._fpn_from_stages(stages)
+
+
+class OfficialHRNetTeacher(nn.Module):
+    """Frozen HRNet / HigherHRNet-style teacher from MMPose (stages F1..F4)."""
+
+    align_mode = "stages"
+    out_channels = 32
+
+    def __init__(
+        self,
+        config_path: str,
+        checkpoint_path: str,
+        device: str = "cpu",
+        width: int = 32,
+    ):
+        super().__init__()
+        self.config_path = _resolve_path(config_path)
+        self.checkpoint_path = _resolve_path(checkpoint_path)
+        if not self.config_path or not os.path.isfile(self.config_path):
+            raise FileNotFoundError(f"task_config not found: {config_path}")
+        if not self.checkpoint_path or not os.path.isfile(self.checkpoint_path):
+            raise FileNotFoundError(f"task_checkpoint not found: {checkpoint_path}")
+
+        from mmpose.apis import init_model
+
+        self.model = init_model(self.config_path, self.checkpoint_path, device=device)
+        freeze_module(self.model)
+        self.width = width
+        self.out_channels = width
+        # projection when codec h channels != stem width
+        self.h_proj = freeze_module(nn.Conv2d(128, width, 1))
+
+    def _preprocess(self, images: torch.Tensor) -> torch.Tensor:
+        return _imagenet_norm_rgb01(images)
+
+    def _stages_from_backbone(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        bb = self.model.backbone if hasattr(self.model, "backbone") else self.model
+        # HRNet forward usually returns multi-resolution list / tensor
+        feats = bb(x)
+        if isinstance(feats, (list, tuple)):
+            # take highest-res stream as F1, then downsample proxies for F2..F4
+            f1 = feats[0] if feats[0].dim() == 4 else feats[0][-1]
+            outs = [f1]
+            cur = f1
+            for i in range(1, 4):
+                if i < len(feats) and isinstance(feats[i], torch.Tensor) and feats[i].dim() == 4:
+                    outs.append(feats[i])
+                    cur = feats[i]
+                else:
+                    cur = F.avg_pool2d(cur, 2, 2)
+                    outs.append(cur)
+            return {f"f{i+1}": outs[i] for i in range(4)}
+        # single tensor: synthesize pyramid
+        f1 = feats
+        outs = [f1]
+        cur = f1
+        for _ in range(3):
+            cur = F.avg_pool2d(cur, 2, 2)
+            outs.append(cur)
+        return {f"f{i+1}": outs[i] for i in range(4)}
+
+    def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]:
+        return self._stages_from_backbone(self._preprocess(images))
+
+    def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]:
+        if h.shape[1] != self.width:
+            if self.h_proj.in_channels != h.shape[1]:
+                self.h_proj = freeze_module(nn.Conv2d(h.shape[1], self.width, 1).to(h.device))
+            h = self.h_proj(h)
+        # Approximate remaining stages with stride-2 pools (HRNet full from-h
+        # needs model-specific stem hooks; pyramid still provides multi-scale D).
+        outs = [h]
+        cur = h
+        for _ in range(3):
+            cur = F.avg_pool2d(cur, 2, 2)
+            outs.append(cur)
+        return {f"f{i+1}": outs[i] for i in range(4)}

flexicm/tasks/swin_teacher.py

@@ -55,12 +55,16 @@ def build_swin_backbone(pretrained: bool = True, swin_variant: str = "base"):
     if key not in _SWIN_TIMM_NAMES:
         raise ValueError(f"Unknown swin_variant={swin_variant!r}; expected one of {list(_SWIN_TIMM_NAMES)}")
 
+    # Training uses 256x256 crops; validation pads to multiples of 256.
+    # Allow variable spatial sizes (do not hard-lock PatchEmbed to 224).
     model = timm.create_model(
         _SWIN_TIMM_NAMES[key],
         pretrained=pretrained,
         features_only=True,
         out_indices=(0, 1, 2, 3),
-        img_size=224,
+        img_size=256,
+        dynamic_img_size=True,
+        strict_img_size=False,
     )
     return model
 

flexicm/utils/codec_test.py

@@ -93,13 +93,13 @@ def test_taic_loader(
         N, _, H, W = images.shape
         num_pixels = N * H * W
 
+        # Keep codec + teacher on the padded grid so Swin always sees a
+        # patch/window-divisible size. bpp still uses the original pixel count.
         x, _ = pad_for_codec(images, divisor=align_divisor, device=device)
         out = model(x)
-        h = crop_feature_to_image(out["h"], (H, W))
-        out["h"] = h
 
-        gt = teacher.gt_features(images)
-        pred = teacher.pred_features(h)
+        gt = teacher.gt_features(x)
+        pred = teacher.pred_features(out["h"])
         stats = criterion(out, pred, gt, num_pixels=num_pixels)
 
         meters["loss"].update(stats["loss"].item(), n=N)
@@ -170,6 +170,7 @@ def test_ctaic_loader(
         N, _, H, W = images.shape
         num_pixels = N * H * W
 
+        # Keep codec + teacher on the padded grid (same rationale as TAIC test).
         x, _ = pad_for_codec(images, divisor=align_divisor, device=device)
         y_b = None
         if use_condition:
@@ -177,11 +178,9 @@ def test_ctaic_loader(
             y_b = base_out["y_hat"]
 
         out = ext_model(x, y_b_hat=y_b, use_condition=use_condition and y_b is not None)
-        h = crop_feature_to_image(out["h"], (H, W))
-        out["h"] = h
 
-        gt = teacher.gt_features(images)
-        pred = teacher.pred_features(h)
+        gt = teacher.gt_features(x)
+        pred = teacher.pred_features(out["h"])
         stats = criterion(out, pred, gt, num_pixels=num_pixels)
 
         meters["loss"].update(stats["loss"].item(), n=N)

scripts/convert_swin_det_ckpt_to_mmdet3.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Convert official Swin-Det Cascade Mask R-CNN ckpt -> mmdet 3.x key layout.
+
+Official source:
+  https://github.com/SwinTransformer/storage/releases/download/v1.0.2/cascade_mask_rcnn_swin_base_patch4_window7.pth
+
+Only the Swin backbone keys need remapping (layers->stages, attn/mlp naming).
+neck / rpn_head / roi_head keys already match mmdet 3 CascadeRCNN.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import torch
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if REPO_ROOT not in sys.path:
+    sys.path.insert(0, REPO_ROOT)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--src",
+        default=os.path.join(REPO_ROOT, "checkpoints/task_networks/detection/model.pth"),
+    )
+    parser.add_argument(
+        "--dst",
+        default=os.path.join(REPO_ROOT, "checkpoints/task_networks/detection/model_mmdet3.pth"),
+    )
+    args = parser.parse_args()
+
+    from mmdet.models.backbones.swin import swin_converter
+
+    raw = torch.load(args.src, map_location="cpu")
+    state = raw["state_dict"] if isinstance(raw, dict) and "state_dict" in raw else raw
+
+    backbone = {}
+    rest = {}
+    for k, v in state.items():
+        if k.startswith("backbone."):
+            backbone[k[len("backbone.") :]] = v
+        else:
+            rest[k] = v
+
+    converted_backbone = swin_converter(backbone)  # adds 'backbone.' prefix
+    out_state = {}
+    out_state.update(converted_backbone)
+    out_state.update(rest)
+
+    os.makedirs(os.path.dirname(args.dst), exist_ok=True)
+    torch.save({"state_dict": out_state, "meta": {"converted_from": args.src}}, args.dst)
+    print(f"Wrote {args.dst}")
+    print(f"  backbone keys: {sum(1 for k in out_state if k.startswith('backbone.'))}")
+    print(f"  other keys:    {sum(1 for k in out_state if not k.startswith('backbone.'))}")
+    print(f"  sample backbone: {[k for k in out_state if k.startswith('backbone.')][:5]}")
+
+
+if __name__ == "__main__":
+    main()

scripts/eval_ctaic.py

@@ -49,7 +49,12 @@ def parse_args(argv):
     given, remaining = parser.parse_known_args(argv)
     cfg_path = given.config if os.path.isabs(given.config) else os.path.join(REPO_ROOT, given.config)
     cfg = load_yaml_config(cfg_path)
-    parser.set_defaults(**cfg)
+    # -c already consumed by the first parse; keep it as a default for the second pass
+    parser.set_defaults(config=cfg_path, **cfg)
+    for action in parser._actions:
+        if "--config" in action.option_strings:
+            action.required = False
+            break
     parser.add_argument("--actual-bpp", action="store_true")
     parser.add_argument("--no-condition", action="store_true")
     parser.add_argument("--with-metrics", action="store_true")
@@ -139,7 +144,14 @@ def main(argv):
     net.load_state_dict(state, strict=False)
     net.eval()
 
-    teacher = build_teacher(ext_task, pretrained_backbone=getattr(args, "pretrained_backbone", True))
+    teacher = build_teacher(
+        ext_task,
+        pretrained_backbone=getattr(args, "pretrained_backbone", True),
+        use_official_teacher=getattr(args, "use_official_teacher", True),
+        task_config=getattr(args, "task_config", None),
+        task_checkpoint=getattr(args, "task_checkpoint", None),
+        device=device,
+    )
     teacher = teacher.to(device).eval()
     criterion = TAICCriterion(lmbda=lmbda, align_mode=align_mode)
     codec_loader = build_codec_loader(args, ext_task, device)

scripts/eval_taic.py

@@ -55,7 +55,12 @@ def parse_args(argv):
     given, remaining = parser.parse_known_args(argv)
     cfg_path = given.config if os.path.isabs(given.config) else os.path.join(REPO_ROOT, given.config)
     cfg = load_yaml_config(cfg_path)
-    parser.set_defaults(**cfg)
+    # -c already consumed by the first parse; keep it as a default for the second pass
+    parser.set_defaults(config=cfg_path, **cfg)
+    for action in parser._actions:
+        if "--config" in action.option_strings:
+            action.required = False
+            break
     parser.add_argument("--actual-bpp", action="store_true")
     parser.add_argument("--with-metrics", action="store_true", help="Run full task-network metrics")
     parser.add_argument("--max-batches", type=int, default=None)
@@ -136,7 +141,14 @@ def main(argv):
     net.eval()
 
     # ---- codec test ----
-    teacher = build_teacher(task, pretrained_backbone=getattr(args, "pretrained_backbone", True))
+    teacher = build_teacher(
+        task,
+        pretrained_backbone=getattr(args, "pretrained_backbone", True),
+        use_official_teacher=getattr(args, "use_official_teacher", True),
+        task_config=getattr(args, "task_config", None),
+        task_checkpoint=getattr(args, "task_checkpoint", None),
+        device=device,
+    )
     teacher = teacher.to(device).eval()
     criterion = TAICCriterion(lmbda=lmbda, align_mode=align_mode)
     codec_loader = build_codec_loader(args, device)

scripts/train_ctaic.py

@@ -35,6 +35,7 @@ from flexicm.data import COCOImageDataset, COCOWholeBodyImageDataset, build_trai
 from flexicm.models import CTAIC, TAIC
 from flexicm.tasks import TASK_META, build_teacher
 from flexicm.tasks.losses import TAICCriterion
+from flexicm.utils.alignment import Alignment
 from flexicm.utils.train_utils import (
     AverageMeter,
     CustomDataParallel,
@@ -61,7 +62,12 @@ def parse_args(argv):
     parser.add_argument("--name", default=datetime.now().strftime("%Y-%m-%d_%H_%M_%S"))
     given, remaining = parser.parse_known_args(argv)
     cfg = load_yaml_config(given.config)
-    parser.set_defaults(**cfg)
+    # -c/--stage already consumed by the first parse; keep them for the second pass
+    parser.set_defaults(config=given.config, stage=given.stage, **cfg)
+    for action in parser._actions:
+        if "--config" in action.option_strings:
+            action.required = False
+            break
     args = parser.parse_args(remaining)
     args.config = given.config
     args.stage = given.stage
@@ -121,17 +127,20 @@ def train_one_epoch(stage, ext_model, base_model, teacher, loader, optimizer, cr
 
 
 @torch.no_grad()
-def validate(stage, ext_model, base_model, teacher, loader, criterion, device):
+def validate(stage, ext_model, base_model, teacher, loader, criterion, device, align_divisor=256):
     ext_model.eval()
     meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")}
     for images in loader:
         images = images.to(device)
+        # Pad to codec/Swin-friendly size (same as train_taic.validate).
+        align = Alignment(divisor=align_divisor, mode="pad", padding_mode="constant").to(device)
+        x = align.align(images)
         if stage == 1:
-            out = ext_model(images, use_condition=False)
+            out = ext_model(x, use_condition=False)
         else:
-            y_b = encode_base_latent(base_model, images)
-            out = ext_model(images, y_b_hat=y_b, use_condition=True)
-        gt = teacher.gt_features(images)
+            y_b = encode_base_latent(base_model, x)
+            out = ext_model(x, y_b_hat=y_b, use_condition=True)
+        gt = teacher.gt_features(x)
         pred = teacher.pred_features(out["h"])
         N, _, H, W = images.shape
         stats = criterion(out, pred, gt, num_pixels=N * H * W)
@@ -198,7 +207,14 @@ def main(argv):
         f"Trainable params: {sum(p.numel() for p in net.parameters() if p.requires_grad)/1e6:.3f}M"
     )
 
-    teacher = build_teacher(ext_task, pretrained_backbone=getattr(args, "pretrained_backbone", True))
+    teacher = build_teacher(
+        ext_task,
+        pretrained_backbone=getattr(args, "pretrained_backbone", True),
+        use_official_teacher=getattr(args, "use_official_teacher", True),
+        task_config=getattr(args, "task_config", None),
+        task_checkpoint=getattr(args, "task_checkpoint", None),
+        device=device,
+    )
     teacher = teacher.to(device).eval()
 
     if args.cuda and torch.cuda.device_count() > 1:

scripts/train_taic.py

@@ -51,7 +51,12 @@ def parse_args(argv):
     parser.add_argument("--name", default=datetime.now().strftime("%Y-%m-%d_%H_%M_%S"))
     given, remaining = parser.parse_known_args(argv)
     cfg = load_yaml_config(given.config)
-    parser.set_defaults(**cfg)
+    # -c already consumed by the first parse; keep it as a default for the second pass
+    parser.set_defaults(config=given.config, **cfg)
+    for action in parser._actions:
+        if "--config" in action.option_strings:
+            action.required = False
+            break
     parser.add_argument("-T", "--TEST", action="store_true")
     args = parser.parse_args(remaining)
     args.config = given.config
@@ -67,12 +72,12 @@ def validate(model, teacher, loader, criterion, device, align_divisor=256):
         align = Alignment(divisor=align_divisor, mode="pad", padding_mode="constant").to(device)
         x = align.align(images)
         out = model(x)
-        # resume spatial pad on h
-        h = align.resume(out["h"])
-        out["h"] = h
+        # Keep teacher + codec features on the aligned (padded) grid so the Swin
+        # backbone always sees a patch/window-divisible size. bpp still uses the
+        # original pixel count below.
         with torch.no_grad():
-            gt = teacher.gt_features(images)
-            pred = teacher.pred_features(h)
+            gt = teacher.gt_features(x)
+            pred = teacher.pred_features(out["h"])
         N, _, H, W = images.shape
         stats = criterion(out, pred, gt, num_pixels=N * H * W)
         for k in meters:
@@ -126,9 +131,9 @@ def main(argv):
     # data
     train_tf = build_train_transform(args.patch_size)
     if task == "pose":
-        train_set = COCOWholeBodyImageDataset(args.dataset_path, "train2017", train_tf)
+        train_set = COCOWholeBodyImageDataset(args.dataset_path, "val2017", train_tf)
     else:
-        train_set = COCOImageDataset(args.dataset_path, "train2017", train_tf)
+        train_set = COCOImageDataset(args.dataset_path, "val2017", train_tf)
     val_split = getattr(args, "val_split", "val2017")
     val_root = os.path.join(args.dataset_path, val_split)
     if not os.path.isdir(val_root):
@@ -163,13 +168,26 @@ def main(argv):
         logging.info(f"Loading base TIC codec from {args.base_codec}")
         state, _ = load_checkpoint_dict(args.base_codec, map_location=device)
         net.load_base_codec(state, strict=False)
-    net.freeze_base_codec()
+    if getattr(args, "freeze_base_codec", True):
+        net.freeze_base_codec()
+        logging.info("Base TIC codec frozen")
+    else:
+        for p in net.parameters():
+            p.requires_grad = True
+        logging.info("Base TIC codec unfrozen; training the entire TAIC model")
     logging.info(
         f"Trainable params: {sum(p.numel() for p in net.parameters() if p.requires_grad)/1e6:.3f}M / "
         f"total {sum(p.numel() for p in net.parameters())/1e6:.3f}M"
     )
 
-    teacher = build_teacher(task, pretrained_backbone=getattr(args, "pretrained_backbone", True))
+    teacher = build_teacher(
+        task,
+        pretrained_backbone=getattr(args, "pretrained_backbone", True),
+        use_official_teacher=getattr(args, "use_official_teacher", True),
+        task_config=getattr(args, "task_config", None),
+        task_checkpoint=getattr(args, "task_checkpoint", None),
+        device=device,
+    )
     teacher = teacher.to(device).eval()
 
     if args.checkpoint:
@@ -192,22 +210,23 @@ def main(argv):
     for epoch in range(args.epochs):
         logging.info(f"===== Epoch {epoch}/{args.epochs} =====")
         train_stats = train_one_epoch(net, teacher, train_loader, optimizer, criterion, device)
-        val_stats = validate(net, teacher, val_loader, criterion, device)
-        logging.info(f"train={train_stats} val={val_stats}")
-        is_best = val_stats["loss"] < best
-        best = min(best, val_stats["loss"])
-        if args.save:
-            save_checkpoint(
-                {
-                    "epoch": epoch,
-                    "state_dict": net.module.state_dict() if hasattr(net, "module") else net.state_dict(),
-                    "optimizer": optimizer.state_dict(),
-                    "loss": val_stats["loss"],
-                    "args": vars(args),
-                },
-                is_best,
-                out_dir,
-            )
+        if epoch % 10 == 0:
+            val_stats = validate(net, teacher, val_loader, criterion, device)
+            logging.info(f"train={train_stats} val={val_stats}")
+            is_best = val_stats["loss"] < best
+            best = min(best, val_stats["loss"])
+            if args.save:
+                save_checkpoint(
+                    {
+                        "epoch": epoch,
+                        "state_dict": net.module.state_dict() if hasattr(net, "module") else net.state_dict(),
+                        "optimizer": optimizer.state_dict(),
+                        "loss": val_stats["loss"],
+                        "args": vars(args),
+                    },
+                    is_best,
+                    out_dir,
+                )
 
 
 if __name__ == "__main__":