[Features] : add flow1d correlation and correlation lookup

configs/_base_/datasets/flyingthings3d_flow1d_400x720.py

@@ -0,0 +1,106 @@
+train_dataset_type = 'FlyingThings3D'
+train_data_root = 'data/flyingthings3d'
+test_dataset_type = 'Sintel'
+test_data_root = 'data/Sintel'
+
+img_norm_cfg = dict(
+    mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], to_rgb=False)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(
+        type='ColorJitter',
+        asymmetric_prob=0.2,
+        brightness=0.4,
+        contrast=0.4,
+        saturation=0.4,
+        hue=0.5 / 3.14),
+    dict(type='Erase', prob=0.5, bounds=[50, 100], max_num=3),
+    dict(
+        type='SpacialTransform',
+        spacial_prob=0.8,
+        stretch_prob=0.8,
+        crop_size=(400, 720),
+        min_scale=-0.4,
+        max_scale=0.8,
+        max_stretch=0.2),
+    dict(type='RandomCrop', crop_size=(400, 720)),
+    dict(type='RandomFlip', prob=0.5, direction='horizontal'),
+    dict(type='RandomFlip', prob=0.1, direction='vertical'),
+    dict(type='Validation', max_flow=1000.),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='DefaultFormatBundle'),
+    dict(
+        type='Collect',
+        keys=['imgs', 'flow_gt', 'valid'],
+        meta_keys=[
+            'filename1', 'filename2', 'ori_filename1', 'ori_filename2',
+            'filename_flow', 'ori_filename_flow', 'ori_shape', 'img_shape',
+            'erase_bounds', 'erase_num', 'scale_factor'
+        ])
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(type='InputPad', exponent=3),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='TestFormatBundle'),
+    dict(
+        type='Collect',
+        keys=['imgs'],
+        meta_keys=[
+            'flow_gt', 'filename1', 'filename2', 'ori_filename1',
+            'ori_filename2', 'ori_shape', 'img_shape', 'img_norm_cfg',
+            'scale_factor', 'pad_shape', 'pad'
+        ])
+]
+
+train_dataset_cleanpass = dict(
+    type=train_dataset_type,
+    data_root=train_data_root,
+    pipeline=train_pipeline,
+    test_mode=False,
+    pass_style='clean',
+    scene='left')
+train_dataset_finalpass = dict(
+    type=train_dataset_type,
+    data_root=train_data_root,
+    pipeline=train_pipeline,
+    test_mode=False,
+    pass_style='final',
+    scene='left')
+test_data_cleanpass = dict(
+    type=test_dataset_type,
+    data_root=test_data_root,
+    pipeline=test_pipeline,
+    test_mode=True,
+    pass_style='clean')
+test_data_finalpass = dict(
+    type=test_dataset_type,
+    data_root=test_data_root,
+    pipeline=test_pipeline,
+    test_mode=True,
+    pass_style='final')
+
+data = dict(
+    train_dataloader=dict(
+        samples_per_gpu=2,
+        workers_per_gpu=2,
+        shuffle=False,
+        drop_last=True,
+        persistent_workers=True),
+    val_dataloader=dict(
+        samples_per_gpu=1,
+        workers_per_gpu=2,
+        shuffle=False,
+        persistent_workers=True),
+    test_dataloader=dict(samples_per_gpu=1, workers_per_gpu=2, shuffle=False),
+    train=[train_dataset_cleanpass, train_dataset_finalpass],
+    val=dict(
+        type='ConcatDataset',
+        datasets=[test_data_cleanpass, test_data_finalpass],
+        separate_eval=True),
+    test=dict(
+        type='ConcatDataset',
+        datasets=[test_data_cleanpass, test_data_finalpass],
+        separate_eval=True))

configs/_base_/models/flow1d.py

@@ -0,0 +1,46 @@
+model = dict(
+    type='Flow1D',
+    radius=32,
+    cxt_channels=128,
+    h_channels=128,
+    encoder=dict(
+        type='RAFTEncoder',
+        in_channels=3,
+        out_channels=256,
+        net_type='Basic',
+        norm_cfg=dict(type='IN'),
+        init_cfg=[
+            dict(
+                type='Kaiming',
+                layer=['Conv2d'],
+                mode='fan_out',
+                nonlinearity='relu'),
+            dict(type='Constant', layer=['InstanceNorm2d'], val=1, bias=0)
+        ]),
+    cxt_encoder=dict(
+        type='RAFTEncoder',
+        in_channels=3,
+        out_channels=256,
+        net_type='Basic',
+        norm_cfg=dict(type='SyncBN'),
+        init_cfg=[
+            dict(
+                type='Kaiming',
+                layer=['Conv2d'],
+                mode='fan_out',
+                nonlinearity='relu'),
+            dict(type='Constant', layer=['SyncBatchNorm2d'], val=1, bias=0)
+        ]),
+    decoder=dict(
+        type='Flow1DDecoder',
+        net_type='Basic',
+        radius=32,
+        iters=12,
+        corr_op_cfg=dict(type='CorrLookupFlow1D', align_corners=True),
+        gru_type='SeqConv',
+        flow_loss=dict(type='SequenceLoss'),
+        act_cfg=dict(type='ReLU')),
+    freeze_bn=False,
+    train_cfg=dict(),
+    test_cfg=dict(),
+)

configs/flow1d/flow1d_8x2_100k_flyingchairs_368x496.py

@@ -0,0 +1,24 @@
+_base_ = [
+    '../_base_/models/flow1d.py',
+    '../_base_/datasets/flyingchairs_raft_368x496.py',
+    '../_base_/default_runtime.py'
+]
+
+optimizer = dict(
+    type='AdamW',
+    lr=0.0004,
+    betas=(0.9, 0.999),
+    eps=1e-08,
+    weight_decay=0.0001,
+    amsgrad=False)
+optimizer_config = dict(grad_clip=dict(max_norm=1.))
+lr_config = dict(
+    policy='OneCycle',
+    max_lr=0.0004,
+    total_steps=100100,
+    pct_start=0.05,
+    anneal_strategy='linear')
+
+runner = dict(type='IterBasedRunner', max_iters=100000)
+checkpoint_config = dict(by_epoch=False, interval=10000)
+evaluation = dict(interval=10000, metric='EPE')

configs/flow1d/flow1d_8x2_100k_flyingthings3d_400x720.py

@@ -0,0 +1,29 @@
+_base_ = [
+    '../_base_/models/flow1d.py',
+    '../_base_/datasets/flyingthings3d_raft_400x720.py',
+    '../_base_/default_runtime.py'
+]
+
+model = dict(freeze_bn=True, test_cfg=dict(iters=32))
+
+optimizer = dict(
+    type='AdamW',
+    lr=0.000125,
+    betas=(0.9, 0.999),
+    eps=1e-08,
+    weight_decay=0.00001,
+    amsgrad=False)
+optimizer_config = dict(grad_clip=dict(max_norm=1.))
+lr_config = dict(
+    policy='OneCycle',
+    max_lr=0.000125,
+    total_steps=100100,
+    pct_start=0.05,
+    anneal_strategy='linear')
+
+runner = dict(type='IterBasedRunner', max_iters=100000)
+checkpoint_config = dict(by_epoch=False, interval=10000)
+evaluation = dict(interval=10000, metric='EPE')
+
+# Train on FlyingChairs and finetune on FlyingThings3D
+load_from = "work-dir/flow1d/flyingchair/latest.pth"

configs/flow1d/flow1d_8x2_100k_flyingthings3d_sintel_368x768.py

@@ -0,0 +1,41 @@
+_base_ = [
+    '../_base_/models/flow1d.py',
+    '../_base_/datasets/sintel_cleanx100_sintel_fianlx100_flyingthings3d_raft_368x768.py',  # noqa
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    decoder=dict(
+        type='Flow1DDecoder',
+        net_type='Basic',
+        radius=32,
+        iters=12,
+        corr_op_cfg=dict(type='CorrLookupFlow1D', align_corners=True),
+        gru_type='SeqConv',
+        flow_loss=dict(type='SequenceLoss', gamma=0.85),
+        act_cfg=dict(type='ReLU')),
+    freeze_bn=True,
+    test_cfg=dict(iters=32))
+
+optimizer = dict(
+    type='AdamW',
+    lr=0.000125,
+    betas=(0.9, 0.999),
+    eps=1e-08,
+    weight_decay=0.00001,
+    amsgrad=False)
+optimizer_config = dict(grad_clip=dict(max_norm=1.))
+lr_config = dict(
+    policy='OneCycle',
+    max_lr=0.000125,
+    total_steps=100100,
+    pct_start=0.05,
+    anneal_strategy='linear')
+
+runner = dict(type='IterBasedRunner', max_iters=100000)
+checkpoint_config = dict(by_epoch=False, interval=10000)
+evaluation = dict(interval=10000, metric='EPE')
+
+# Train on FlyingChairs and FlyingThings3D, and finetune on FlyingThings3D
+# and Sintel
+load_from = "work-dir/flow1d/flyingthingd/latest.pth"

configs/flow1d/flow1d_8x2_100k_mixed_368x768.py

@@ -0,0 +1,42 @@
+_base_ = [
+    '../_base_/models/flow1d.py',
+    '../_base_/datasets/sintel_cleanx100_sintel_fianlx100_kitti2015x200_hd1kx5_flyingthings3d_raft_384x768.py',  # noqa
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    decoder=dict(
+        type='Flow1DDecoder',
+        net_type='Basic',
+        radius=32,
+        iters=12,
+        corr_op_cfg=dict(type='CorrLookupFlow1D', align_corners=True),
+        gru_type='SeqConv',
+        flow_loss=dict(type='SequenceLoss', gamma=0.85),
+        act_cfg=dict(type='ReLU')),
+    freeze_bn=True,
+    test_cfg=dict(iters=32))
+
+optimizer = dict(
+    type='AdamW',
+    lr=0.000125,
+    betas=(0.9, 0.999),
+    eps=1e-08,
+    weight_decay=0.00001,
+    amsgrad=False)
+optimizer_config = dict(grad_clip=dict(max_norm=1.))
+lr_config = dict(
+    policy='OneCycle',
+    max_lr=0.000125,
+    total_steps=100100,
+    pct_start=0.05,
+    anneal_strategy='linear')
+
+runner = dict(type='IterBasedRunner', max_iters=100000)
+checkpoint_config = dict(by_epoch=False, interval=10000)
+evaluation = dict(interval=10000, metric='EPE')
+
+# Train on FlyingChairs and FlyingThings3D, and finetune on
+# and Sintel, KITTI2015 and HD1K
+load_from = "work-dir/flow1d/sintel/latest.pth"
+

mmflow/models/decoders/__init__.py

@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .context_net import ContextNet
+from .flow1d_decoder import Flow1DDecoder
 from .flownet_decoder import FlowNetCDecoder, FlowNetSDecoder
 from .gma_decoder import GMADecoder
 from .irr_refine import FlowRefine, OccRefine, OccShuffleUpsample
@@ -13,5 +14,5 @@
     'FlowNetCDecoder', 'FlowNetSDecoder', 'PWCNetDecoder',
     'MaskFlowNetSDecoder', 'NetE', 'ContextNet', 'RAFTDecoder', 'FlowRefine',
     'OccRefine', 'OccShuffleUpsample', 'IRRPWCDecoder', 'MaskFlowNetDecoder',
-    'GMADecoder'
+    'GMADecoder', 'Flow1DDecoder'
 ]