Mask2Former doesn't converge and worse then Mask-RCNN

[Robotatron]

Training Mask2Former with the Detectron 2 framework provides higher mAP then Mask-RCNN as expected.

However, doing the same experiment within mmdet with the default config does not converge with different learning rates. Mask-RCNN works great (blue graph):

I tried different learning rates:

Any help is appreciated! Mask2Former should outperform Mask-RCNN by a large margin, however it doesnt converge in mmdet, while working perfectly with Detectron 2.

Full config:

base = ['./mask2former_r50_lsj_8x2_50e_coco-panoptic.py']
pretrained = "/home/jupyter/notebooks/Unity/mmdetection/weights/mask2former_swin-t-p4-w7-224_lsj_8x2_50e_coco_20220508_091649-4a943037.pth"
load_from = pretrained

num_things_classes = 7
num_stuff_classes = 0
num_classes = num_things_classes + num_stuff_classes
model = dict(
panoptic_head=dict(
num_things_classes=num_things_classes,
num_stuff_classes=num_stuff_classes,
loss_cls=dict(class_weight=[1.0] * num_classes + [0.1])),
panoptic_fusion_head=dict(
num_things_classes=num_things_classes,
num_stuff_classes=num_stuff_classes),
test_cfg=dict(panoptic_on=False),
)

#------------------------------------------------------------------------------------------------------------------------

data settings

#------------------------------------------------------------------------------------------------------------------------
image_size = (1024, 1024)
img_norm_cfg = dict(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
pad_cfg = dict(img=(128, 128, 128), masks=0, seg=255)
train_pipeline = [
dict(type='LoadImageFromFile', to_float32=True),
dict(type='LoadAnnotations', with_bbox=True, with_mask=True),
dict(type='RandomFlip', flip_ratio=0.5),
# large scale jittering
dict(
type='Resize',
img_scale=image_size,
ratio_range=(0.1, 2.0),
multiscale_mode='range',
keep_ratio=True),
dict(
type='RandomCrop',
crop_size=image_size,
crop_type='absolute',
recompute_bbox=True,
allow_negative_crop=True),
dict(type='FilterAnnotations', min_gt_bbox_wh=(1e-5, 1e-5), by_mask=True),
dict(type='Pad', size=image_size, pad_val=pad_cfg),
#dict(type='Normalize', **img_norm_cfg),
dict(type='DefaultFormatBundle', img_to_float=True),
dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
]
test_pipeline = [
dict(type='LoadImageFromFile'),
dict(
type='MultiScaleFlipAug',
img_scale=image_size,#(1333, 800),
flip=False,
transforms=[
dict(type='Resize', keep_ratio=True),
dict(type='RandomFlip'),
dict(type='Pad', size_divisor=32, pad_val=pad_cfg),
dict(type='Normalize', **img_norm_cfg),
dict(type='ImageToTensor', keys=['img']),
dict(type='Collect', keys=['img']),
])
]

dataset

dataset_type = 'CocoDataset'
data_root = '/home/jupyter/data/Unity/05-Segm-4cars-4k/converted/'
img_prefix = data_root + "images"
data = dict(
delete=True,
samples_per_gpu=1, #batch size bs
workers_per_gpu=1,
train=dict(
type=dataset_type,
ann_file=data_root + 'annotations/train.json',
img_prefix=img_prefix,
pipeline=train_pipeline,
classes = ("hood", "Door-Front", "roof", "tyre", "Door-Rear", "Bumper-Front", "Bumper-Rear")
),
val=dict(
type=dataset_type,
ann_file=data_root + 'annotations/test.json',
img_prefix=img_prefix,
pipeline=test_pipeline,
classes = ("hood", "Door-Front", "roof", "tyre", "Door-Rear", "Bumper-Front", "Bumper-Rear")
),
test=dict(
type=dataset_type,
ann_file=data_root + 'annotations/test.json',
img_prefix=img_prefix,
pipeline=test_pipeline,
classes = ("hood", "Door-Front", "roof", "tyre", "Door-Rear", "Bumper-Front", "Bumper-Rear")
)
)
#------------------------------------------------------------------------------------------------------------------------

embed_multi = dict(lr_mult=1.0, decay_mult=0.0)

optimizer

optimizer = dict(
type='AdamW',
lr=0.0001, #0.001
weight_decay=0.05,
eps=1e-8,
betas=(0.9, 0.999),
paramwise_cfg=dict(
custom_keys={
'backbone': dict(lr_mult=0.1, decay_mult=1.0),
'query_embed': embed_multi,
'query_feat': embed_multi,
'level_embed': embed_multi,
},
norm_decay_mult=0.0))

max_iters = 100_000
runner = dict(type='IterBasedRunner', max_iters=max_iters)

work_dir = "work_dirs/10-Unity-04-Mask2Former"
log_config = dict(
interval=100,
hooks=[
dict(type='TextLoggerHook', by_epoch=False),
dict(type='TensorboardLoggerHook', by_epoch=False, log_dir=work_dir)
])
interval = 1000
workflow = [('train', interval)]
checkpoint_config = dict(
by_epoch=False, interval=interval, save_last=True, max_keep_ckpts=3)

Before 365001th iteration, we do evaluation every 5000 iterations.

After 365000th iteration, we do evaluation every 368750 iterations,

which means that we do evaluation at the end of training.

dynamic_intervals = [(max_iters // interval * interval + 1, max_iters)]
evaluation = dict(
interval=interval,
dynamic_intervals=dynamic_intervals,
metric=['segm'])

[ccomkhj]

I also faced this problem.
My reasoning is that transformer is much more sensitive than Mask-RCNN in training.
Depends on the lr_config and learning rate, loss were significantly different.

[Robotatron]

@ccomkhj

I also faced this problem. My reasoning is that transformer is much more sensitive than Mask-RCNN in training. Depends on the lr_config and learning rate, loss were significantly different.

Were you able to "fix" it?
The default learning rate is 0.0001 but the default config was used with 8x GPUs. If I am training on a single GPU, should I divide the learning rate by two?
Same for the image size, default is 8 or 16, I am using batch size of 2.

[ccomkhj]

I made "meaningful" progress setting up the step policy of lr_config by trying multiple step values.
(More than 4 steps)
So, train loss converged, but the final performance of MaskRCNN is still better than Mask2Former.
I believe it is due to the shortage of datasets. (Transformer needs far more datasets than CNN based network.)