How to use mmselfsup to train with coco dataset?

[fangyixiao18]

In old version of MMSelfSup, it is not convenient to use mmselfsup to train algorithms with other datasets because of the different pipeline of the datasets.

But now, with the v1.0.0rc version released, the users are able to train self-supervised learning algorithms with other datasets conveniently, such as coco.

Here is an example to train MAE with coco dataset.

1. Make sure that you have installed MMDetection with correct version, for example, MMSelfSup v1.0.0rc2 with MMDetection v3.0.0rc1.
2. Modify your config.
   The original config of MAE:

# dataset settings
custom_imports = dict(imports='mmcls.datasets', allow_failed_imports=False)
dataset_type = 'mmcls.ImageNet'
data_root = 'data/imagenet/'
file_client_args = dict(backend='disk')

train_pipeline = [
    dict(type='LoadImageFromFile', file_client_args=file_client_args),
    dict(
        type='RandomResizedCrop',
        size=224,
        scale=(0.2, 1.0),
        backend='pillow',
        interpolation='bicubic'),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PackSelfSupInputs', meta_keys=['img_path'])
]

train_dataloader = dict(
    batch_size=128,
    num_workers=8,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    collate_fn=dict(type='default_collate'),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        ann_file='meta/train.txt',
        data_prefix=dict(img_path='train/'),
        pipeline=train_pipeline))

Create a new config called coco_mae.py and modify the original config of MAE with coco dataset

# dataset settings
dataset_type = 'mmdet.CocoDataset'
data_root = 'data/coco/'
file_client_args = dict(backend='disk')

train_pipeline = [
    dict(type='LoadImageFromFile', file_client_args=file_client_args),
    dict(
        type='RandomResizedCrop',
        size=224,
        scale=(0.2, 1.0),
        backend='pillow',
        interpolation='bicubic'),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PackSelfSupInputs', meta_keys=['img_path'])
]

train_dataloader = dict(
    batch_size=128,
    num_workers=8,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    collate_fn=dict(type='default_collate'),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        ann_file='annotations/instances_train2017.json',
        data_prefix=dict(img='train2017/'),
        pipeline=train_pipeline))

The main different is the dataset_type, data_root and the key in train_dataloader['dataset']

3. In the algorithm config, like mae_vit-base-p16_8xb512-amp-coslr-300e_coco.py, change the _base_ file.

_base_ = [
    '../_base_/models/mae_vit-base-p16.py',
    '../_base_/datasets/coco_mae.py',
    '../_base_/schedules/adamw_coslr-200e_in1k.py',
    '../_base_/default_runtime.py',
]

4. Then start to train!
   Here is a snippet of my training log.
   1 A100 GPU with batchsize=128

2022/10/13 17:35:26 - mmengine - INFO - Checkpoints will be saved to /fangyixiao18/work_dirs/selfsup/mae_vit-base-p16_8xb512-amp-coslr-300e_coco/20221013_test-mae-coco by HardDiskBackend.
2022/10/13 17:36:03 - mmengine - INFO - Epoch(train) [1][100/925]  lr: 1.6653e-06  eta: 1 day, 4:36:48  time: 0.3328  data_time: 0.0053  memory: 18150  loss: 1.2406
2022/10/13 17:36:37 - mmengine - INFO - Epoch(train) [1][200/925]  lr: 3.2868e-06  eta: 1 day, 3:20:40  time: 0.3359  data_time: 0.0052  memory: 18150  loss: 0.9834
2022/10/13 17:37:10 - mmengine - INFO - Epoch(train) [1][300/925]  lr: 4.9083e-06  eta: 1 day, 2:52:55  time: 0.3347  data_time: 0.0052  memory: 18150  loss: 0.8888
2022/10/13 17:37:44 - mmengine - INFO - Epoch(train) [1][400/925]  lr: 6.5298e-06  eta: 1 day, 2:38:41  time: 0.3393  data_time: 0.0067  memory: 18150  loss: 0.8541
2022/10/13 17:38:18 - mmengine - INFO - Epoch(train) [1][500/925]  lr: 8.1513e-06  eta: 1 day, 2:29:33  time: 0.3343  data_time: 0.0052  memory: 18150  loss: 0.8266
2022/10/13 17:38:52 - mmengine - INFO - Epoch(train) [1][600/925]  lr: 9.7728e-06  eta: 1 day, 2:25:42  time: 0.3322  data_time: 0.0050  memory: 18150  loss: 0.8279
2022/10/13 17:39:25 - mmengine - INFO - Epoch(train) [1][700/925]  lr: 1.1394e-05  eta: 1 day, 2:20:07  time: 0.3340  data_time: 0.0051  memory: 18150  loss: 0.7978
2022/10/13 17:40:04 - mmengine - INFO - Epoch(train) [1][800/925]  lr: 1.3016e-05  eta: 1 day, 2:44:36  time: 0.3726  data_time: 0.0056  memory: 18150  loss: 0.7692
2022/10/13 17:40:38 - mmengine - INFO - Epoch(train) [1][900/925]  lr: 1.4637e-05  eta: 1 day, 2:39:44  time: 0.3398  data_time: 0.0055  memory: 18150  loss: 0.7598
2022/10/13 17:40:46 - mmengine - INFO - Exp name: mae_vit-base-p16_8xb512-amp-coslr-300e_coco_20221013_173449
2022/10/13 17:40:46 - mmengine - INFO - Saving checkpoint at 1 epochs

Feel free to comment here !

[fangyixiao18]

As above