How to load LMDB file to speed up SR training?

[Senwang98]

Hi
After creating lmdb file of DIV2K, I want to use lmdb to train RDN x2.
Can you give me some suggestions about changes of RDN config?

[innerlee]

You can search for lmdb datasets in the repo for your need, e.g. https://github.com/open-mmlab/mmediting/blob/23213c839ff2d1907a80d6ea29f13c32a24bb8ef/mmedit/datasets/sr_lmdb_dataset.py#L9

and replace the dataset type in the config with it

[Senwang98]

Ok, I will try it!

[Senwang98]

@ckkelvinchan
I change the config as follow, but it didn't work.

Can you provide the demo config of loading lmbd? I think many other persons may meet the same question!

[ckkelvinchan]

Could you provide a detailed log about how it failed? Moreover, you should use data/DIV2K/DIV2K_train_LR_bicubic_X2_sub.lmdb in lq_folder. Similar for gt_folder.

[Senwang98]

The is my RDN config:

exp_name = 'rdn_x2c64b16_g1_1000k_div2k'

scale = 2
# model settings
model = dict(
    type='BasicRestorer',
    generator=dict(
        type='RDN',
        in_channels=3,
        out_channels=3,
        mid_channels=64,
        num_blocks=16,
        upscale_factor=scale),
    pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'))
# model training and testing settings
train_cfg = None
test_cfg = dict(metrics=['PSNR', 'SSIM'], crop_border=scale)

# dataset settings
train_dataset_type = 'SRAnnotationDataset'
val_dataset_type = 'SRFolderDataset'
train_pipeline = [
    dict(
        type='LoadImageFromFile',
        io_backend='lmdb',
        key='lq',
        flag='color',
        channel_order='rgb'),
    dict(
        type='LoadImageFromFile',
        io_backend='lmdb',
        key='gt',
        flag='color',
        channel_order='rgb'),
    dict(type='RescaleToZeroOne', keys=['lq', 'gt']),
    dict(type='PairedRandomCrop', gt_patch_size=64),
    dict(
        type='Flip', keys=['lq', 'gt'], flip_ratio=0.5,
        direction='horizontal'),
    dict(type='Flip', keys=['lq', 'gt'], flip_ratio=0.5, direction='vertical'),
    dict(type='RandomTransposeHW', keys=['lq', 'gt'], transpose_ratio=0.5),
    dict(type='Collect', keys=['lq', 'gt'], meta_keys=['lq_path', 'gt_path']),
    dict(type='ImageToTensor', keys=['lq', 'gt'])
]

test_pipeline = [
    dict(
        type='LoadImageFromFile',
        io_backend='disk',
        key='lq',
        flag='color',
        channel_order='rgb'),
    dict(
        type='LoadImageFromFile',
        io_backend='disk',
        key='gt',
        flag='color',
        channel_order='rgb'),
    dict(type='RescaleToZeroOne', keys=['lq', 'gt']),
    dict(type='Collect', keys=['lq', 'gt'], meta_keys=['lq_path', 'gt_path']),
    dict(type='ImageToTensor', keys=['lq', 'gt'])
]

data = dict(
    workers_per_gpu=1,
    train_dataloader=dict(samples_per_gpu=16, drop_last=True),
    val_dataloader=dict(samples_per_gpu=1),
    test_dataloader=dict(samples_per_gpu=1),
    train=dict(
        type='RepeatDataset',
        times=1000,
        dataset=dict(
            type=train_dataset_type,
            # lq_folder='data/DIV2K/DIV2K_train_LR_bicubic/X2_sub',
            # gt_folder='data/DIV2K/DIV2K_train_HR_sub',
            lq_folder='data/DIV2K/DIV2K_train_LR_bicubic_X2_sub.lmdb',
            gt_folder='data/DIV2K/DIV2K_train_HR_sub.lmdb',
            ann_file='data/DIV2K/meta_info_DIV2K800sub_GT.txt',
            pipeline=train_pipeline,
            scale=scale)),
    val=dict(
        type=val_dataset_type,
        lq_folder='data/val_set5/Set5_bicLRx2',
        gt_folder='data/val_set5/Set5_mod12',
        pipeline=test_pipeline,
        scale=scale,
        filename_tmpl='{}'),
    test=dict(
        type=val_dataset_type,
        lq_folder='data/val_set5/Set5_bicLRx2',
        gt_folder='data/val_set5/Set5_mod12',
        pipeline=test_pipeline,
        scale=scale,
        filename_tmpl='{}'))

# optimizer
optimizers = dict(generator=dict(type='Adam', lr=1e-4, betas=(0.9, 0.999)))

# learning policy
total_iters = 1000000
lr_config = dict(
    policy='Step',
    by_epoch=False,
    step=[200000, 400000, 600000, 800000],
    gamma=0.5)

checkpoint_config = dict(interval=5000, save_optimizer=True, by_epoch=False)
evaluation = dict(interval=1000, save_image=True, gpu_collect=True)
log_config = dict(
    interval=100, hooks=[dict(type='TextLoggerHook', by_epoch=False)])
visual_config = None

# runtime settings
dist_params = dict(backend='nccl')
log_level = 'INFO'
work_dir = f'./work_dirs/{exp_name}'
load_from = None
resume_from = None
workflow = [('train', 1)]

I change the train pipeline (disk to lmdb),
The log as follows:

Traceback (most recent call last):
  File "./tools/train.py", line 146, in <module>
    main()
  File "./tools/train.py", line 135, in main
    train_model(
  File "/home/wangsen/ws/mmediting/mmedit/apis/train.py", line 65, in train_model
    _dist_train(
  File "/home/wangsen/ws/mmediting/mmedit/apis/train.py", line 206, in _dist_train
    runner.run(data_loaders, cfg.workflow, cfg.total_iters)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/mmcv/runner/iter_based_runner.py", line 133, in run
    iter_runner(iter_loaders[i], **kwargs)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/mmcv/runner/iter_based_runner.py", line 58, in train
    data_batch = next(data_loader)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/mmcv/runner/iter_based_runner.py", line 32, in __next__
    data = next(self.iter_loader)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 435, in __next__
    data = self._next_data()
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1085, in _next_data
    return self._process_data(data)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1111, in _process_data
    data.reraise()
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/_utils.py", line 428, in reraise
    raise self.exc_type(msg)
TypeError: Caught TypeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 198, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/wangsen/ws/mmediting/mmedit/datasets/dataset_wrappers.py", line 31, in __getitem__
    return self.dataset[idx % self._ori_len]
  File "/home/wangsen/ws/mmediting/mmedit/datasets/base_sr_dataset.py", line 53, in __getitem__
    return self.pipeline(results)
  File "/home/wangsen/ws/mmediting/mmedit/datasets/pipelines/compose.py", line 42, in __call__
    data = t(data)
  File "/home/wangsen/ws/mmediting/mmedit/datasets/pipelines/loading.py", line 63, in __call__
    self.file_client = FileClient(self.io_backend, **self.kwargs)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/mmcv/fileio/file_client.py", line 236, in __init__
    self.client = self._backends[backend](**kwargs)
TypeError: __init__() missing 1 required positional argument: 'db_path'

Traceback (most recent call last):
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/distributed/launch.py", line 260, in <module>
    main()
  File "/home/wangsen/miniconda3/envs/pt1.7/lib/python3.8/site-packages/torch/distributed/launch.py", line 255, in main
    raise subprocess.CalledProcessError(returncode=process.returncode,
subprocess.CalledProcessError: Command '['/home/wangsen/miniconda3/envs/pt1.7/bin/python', '-u', './tools/train.py', '--local_rank=0', 'configs/restorers/rdn/rdn_x2c64b16_g1_1000k_div2k.py', '--launcher', 'pytorch']' returned non-zero exit status 1.

[Senwang98]

If I don't change train_pipeline and keep disk unchanged,
The error log as follows:
FileNotFoundError: [Errno 2] No such file or directory: 'data/DIV2K/DIV2K_train_LR_bicubic_X2_sub.lmdb/0430_s037.png'

[ckkelvinchan]

Try

    dict(
        type='LoadImageFromFile',
        io_backend='lmdb',
        db_path=''data/DIV2K/DIV2K_train_LR_bicubic_X2_sub.lmdb',
        key='lq',
        flag='color',
        channel_order='rgb'),
    dict(
        type='LoadImageFromFile',
        io_backend='lmdb',
        db_path=''data/DIV2K/DIV2K_train_HR_sub.lmdb',
        key='gt',
        flag='color',
        channel_order='rgb'),

and see if it works?

[ckkelvinchan]

And you may need SRLmdbDataset as discussed above

[Senwang98]

@ckkelvinchan
Thanks, it works now!
So, maybe some tips need to be added into the SR instruction.md, this may save much time.

[innerlee]

Hi @Senwang98 would you like to contribute to the documentation, so that later users would benefit from it?

[Senwang98]

@innerlee
Yes, this is a meaningful job, instructions of using LMDB can be added into doc!
I will add some instructions tonight to benefit other users.

[chenjiachengzzz]

您好，能否提供一个完整的使用lmdb数据进行训练的配置文件qaq！我现在可以使用disk的方式进行训练，但是使用lmdb的方式进行训练就报错了。。。， lr和hr的lmdb文件中都存在meta_info文件，该如何配置ann_file.