[Bug] MultiScaleFlipAug: Normalize: __init__() got an unexpected keyword argument 'keep_ratio'

[sammilei]

Checklist

• I have searched related issues but cannot get the expected help.
• 2. I have read the FAQ documentation but cannot get the expected help.
• 3. The bug has not been fixed in the latest version.

Describe the bug

I got this error when converting an .pth model trained on customized data to .engine.

2023-05-24 16:49:23,909 - mmdeploy - INFO - Start pipeline mmdeploy.apis.pytorch2onnx.torch2onnx in subprocess
load checkpoint from local path: /home/Downloads/best_mIoU_iter_11500.pth
Process Process-2:
Traceback (most recent call last):
  File "/home/.local/lib/python3.8/site-packages/mmcv/utils/registry.py", line 69, in build_from_cfg
    return obj_cls(**args)
TypeError: __init__() got an unexpected keyword argument 'keep_ratio'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/.local/lib/python3.8/site-packages/mmcv/utils/registry.py", line 69, in build_from_cfg
    return obj_cls(**args)
  File "/home/Documents/mmsegmentation/mmseg/datasets/pipelines/test_time_aug.py", line 60, in __init__
    self.transforms = Compose(transforms)
  File "/home/Documents/mmsegmentation/mmseg/datasets/pipelines/compose.py", line 23, in __init__
    transform = build_from_cfg(transform, PIPELINES)
  File "/home/.local/lib/python3.8/site-packages/mmcv/utils/registry.py", line 72, in build_from_cfg
    raise type(e)(f'{obj_cls.__name__}: {e}')
TypeError: Normalize: __init__() got an unexpected keyword argument 'keep_ratio'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/usr/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/Documents/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 107, in __call__
    ret = func(*args, **kwargs)
  File "/home/Documents/mmdeploy/mmdeploy/apis/pytorch2onnx.py", line 63, in torch2onnx
    data, model_inputs = task_processor.create_input(img, input_shape)
  File "/home/Documents/mmdeploy/mmdeploy/codebase/mmseg/deploy/segmentation.py", line 121, in create_input
    test_pipeline = Compose(cfg.data.test.pipeline)
  File "/home/Documents/mmsegmentation/mmseg/datasets/pipelines/compose.py", line 23, in __init__
    transform = build_from_cfg(transform, PIPELINES)
  File "/home/.local/lib/python3.8/site-packages/mmcv/utils/registry.py", line 72, in build_from_cfg
    raise type(e)(f'{obj_cls.__name__}: {e}')
TypeError: MultiScaleFlipAug: Normalize: __init__() got an unexpected keyword argument 'keep_ratio'
2023-05-24 16:49:32,495 - mmdeploy - ERROR - `mmdeploy.apis.pytorch2onnx.torch2onnx` with Call id: 0 failed. exit.

Reproduction

ubuntu:~/Documents/mmdeploy$ python3 tools/deploy.py configs/mmseg/segmentation_tensorrt_static-512x512.py ~/Downloads/segmenter.py ~/Downloads/best_mIoU_iter_11500.pth img.jpg --device cuda:0

This machine was the same one in #2098

The config file segmenter.py is

checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_small_p16_384_20220308-410f6037.pth'
backbone_norm_cfg = dict(type='LN', eps=1e-06, requires_grad=True)
model = dict(
    type='EncoderDecoder',
    pretrained=
    'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_small_p16_384_20220308-410f6037.pth',
    backbone=dict(
        type='VisionTransformer',
        img_size=(512, 512),
        patch_size=16,
        in_channels=3,
        embed_dims=384,
        num_layers=12,
        num_heads=6,
        drop_path_rate=0.1,
        attn_drop_rate=0.0,
        drop_rate=0.0,
        final_norm=True,
        norm_cfg=dict(type='LN', eps=1e-06, requires_grad=True),
        with_cls_token=True,
        interpolate_mode='bicubic'),
    decode_head=dict(
        type='SegmenterMaskTransformerHead',
        in_channels=384,
        channels=384,
        num_classes=6,
        num_layers=2,
        num_heads=6,
        embed_dims=384,
        dropout_ratio=0.0,
        loss_decode=dict(
            type='CrossEntropyLoss',
            use_sigmoid=False,
            loss_weight=1.0,
            avg_non_ignore=True,
            class_weight=[0.624, 0.966, 0.772, 0.821, 0.872, 0.945]),
        ignore_index=255),
    test_cfg=dict(mode='slide', crop_size=(512, 512), stride=(480, 480)))
log_config = dict(
    interval=50, hooks=[dict(type='TextLoggerHook', by_epoch=False)])
dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]
cudnn_benchmark = True
optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0)
optimizer_config = dict()
lr_config = dict(policy='poly', power=0.9, min_lr=0.0001, by_epoch=False)
runner = dict(type='IterBasedRunner', max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=2000)
evaluation = dict(interval=500, metric='mIoU', pre_eval=True, save_best='mIoU')
dataset_type = 'CustomDataset'
data_root = '/training_set_v17'
classes = ['ground', 'tree', 'sky', 'trail', 'rock', 'bush']
class_weights = [0.624, 0.966, 0.772, 0.821, 0.872, 0.945]
test_palette = [(255, 218, 165), (137, 189, 134), (195, 243, 255),
                (255, 255, 186), (199, 206, 234), (226, 240, 203)]
img_norm_cfg = dict(
    mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], to_rgb=True)
crop_size = (512, 512)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=False),
    dict(type='Resize', img_scale=(1440, 900), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=(512, 512), cat_max_ratio=0.5),
    dict(type='RandomFlip', prob=0.5),
    dict(
        type='Normalize',
        mean=[127.5, 127.5, 127.5],
        std=[127.5, 127.5, 127.5],
        to_rgb=True),
    dict(type='Pad', size=(512, 512), pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg'])
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(1440, 900),
        flip=False,
        transforms=[
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_rgb=True),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img'])
        ])
]
data = dict(
    samples_per_gpu=4,
    workers_per_gpu=4,
    train=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        '/training_set_v17',
        img_dir='images/train',
        ann_dir='labels/train',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(type='LoadAnnotations', reduce_zero_label=False),
            dict(type='Resize', img_scale=(1440, 900), ratio_range=(0.5, 2.0)),
            dict(type='RandomCrop', crop_size=(512, 512), cat_max_ratio=0.5),
            dict(type='RandomFlip', prob=0.5),
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_rgb=True),
            dict(type='Pad', size=(512, 512), pad_val=0, seg_pad_val=255),
            dict(type='DefaultFormatBundle'),
            dict(type='Collect', keys=['img', 'gt_semantic_seg'])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush']),
    val=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        '/training_set_v17',
        img_dir='images/val',
        ann_dir='labels/val',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1440, 900),
                flip=False,
                transforms=[
                    dict(
                        type='Normalize',
                        mean=[127.5, 127.5, 127.5],
                        std=[127.5, 127.5, 127.5],
                        to_rgb=True),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush']),
    test=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        '/training_set_v17',
        img_dir='images/val',
        ann_dir='labels/val',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1440, 900),
                flip=False,
                transforms=[
                    dict(
                        type='Normalize',
                        mean=[127.5, 127.5, 127.5],
                        std=[127.5, 127.5, 127.5],
                        to_rgb=True),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush'],
        palette=[(255, 218, 165), (137, 189, 134), (195, 243, 255),
                 (255, 255, 186), (199, 206, 234), (226, 240, 203)]))
work_dir = './work_dirs/segmenter'
gpu_ids = [0]
auto_resume = False

Environment

- GCC 9.4
  - C++ Version: 201402
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - LAPACK is enabled (usually provided by MKL)
  - CPU capability usage: NO AVX
  - CUDA Runtime 11.4
  - NVCC architecture flags: -gencode;arch=compute_62,code=sm_62;-gencode;arch=compute_72,code=sm_72;-gencode;arch=compute_87,code=sm_87
  - CuDNN 8.4.1
    - Built with CuDNN 8.3.2

mmdeploy == 0.6.0
mmcv-full == 1.6.0
mmseg == 0.24.1

Error traceback

No response

[irexyc]

Try change your test_pipeline to

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(1440, 900),
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_rgb=True),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img'])
        ])
]

[sammilei]

Thanks @irexyc it might work but now I got the error for the processor:

Process Process-2:
Traceback (most recent call last):
  File "/usr/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/usr/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/Documents/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 107, in __call__
    ret = func(*args, **kwargs)
  File "/home/Documents/mmdeploy/mmdeploy/apis/pytorch2onnx.py", line 87, in torch2onnx
    export(
  File "/home/Documents/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 356, in _wrap
    return self.call_function(func_name_, *args, **kwargs)
  File "/home/Documents/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 326, in call_function
    return self.call_function_local(func_name, *args, **kwargs)
  File "/home/Documents/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 275, in call_function_local
    return pipe_caller(*args, **kwargs)
  File "/home/Documents/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 107, in __call__
    ret = func(, **kwargs)
  File "/home/gvsc/Documents/mmdeploy/mmdeploy/apis/onnx/export.py", line 122, in export
    torch.onnx.export(
  File "/home/.local/lib/python3.8/site-packages/torch/onnx/__init__.py", line 319, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/home/.local/lib/python3.8/site-packages/torch/onnx/utils.py", line 113, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/home/.local/lib/python3.8/site-packages/torch/onnx/utils.py", line 734, in _export
    params_dict = torch._C._jit_pass_onnx_deduplicate_initializers(graph, params_dict,
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument other in method wrapper__equal)
2023-05-25 09:27:17,328 - mmdeploy - ERROR - `mmdeploy.apis.pytorch2onnx.torch2onnx` with Call id: 0 failed. exit.

[irexyc]

My last comment is not right. As you didn't use test_pipeline, you should change data.test.pipeline to

             dict(type='LoadImageFromFile'),
             dict(
                 type='MultiScaleFlipAug',
                 img_scale=(1440, 900),
                 flip=False,
                 transforms=[
                     dict(type='Resize', keep_ratio=True),
                     dict(
                         type='Normalize',
                         mean=[127.5, 127.5, 127.5],
                         std=[127.5, 127.5, 127.5],
                         to_rgb=True),
                     dict(type='ImageToTensor', keys=['img']),
                     dict(type='Collect', keys=['img'])
                 ]
            )

After that, I can convert the model successful.

Have you modified the mmseg repo? I see in #2098, you can convert segmenter_vit-s_linear_8x1_512x512_160k_ade20k.py successful.

[sammilei]

My last comment is not right. As you didn't use test_pipeline, you should change data.test.pipeline to

             dict(type='LoadImageFromFile'),
             dict(
                 type='MultiScaleFlipAug',
                 img_scale=(1440, 900),
                 flip=False,
                 transforms=[
                     dict(type='Resize', keep_ratio=True),
                     dict(
                         type='Normalize',
                         mean=[127.5, 127.5, 127.5],
                         std=[127.5, 127.5, 127.5],
                         to_rgb=True),
                     dict(type='ImageToTensor', keys=['img']),
                     dict(type='Collect', keys=['img'])
                 ]
            )

After that, I can convert the model successful.

Have you modified the mmseg repo? I see in #2098, you can convert segmenter_vit-s_linear_8x1_512x512_160k_ade20k.py successful.

@irexyc Thank you for the reply!
I did not modify the mmseg repo. I can convert the segmenter_vit-s_linear_8x1_512x512_160k_ade20k.py but now I am stuck with my own model trained with customized dataset.

I tried the suggestion here. It threw out the original error:

File "/home/.local/lib/python3.8/site-packages/mmcv/utils/registry.py", line 72, in build_from_cfg
    raise type(e)(f'{obj_cls.__name__}: {e}')
TypeError: MultiScaleFlipAug: Normalize: __init__() got an unexpected keyword argument 'keep_ratio'
2023-05-30 09:03:45,150 - mmdeploy - ERROR - `mmdeploy.apis.pytorch2onnx.torch2onnx` with Call id: 0 failed. exit.

The config py I am using with the modification on test pipeline:

checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_small_p16_384_20220308-410f6037.pth'
backbone_norm_cfg = dict(type='LN', eps=1e-06, requires_grad=True)
model = dict(
    type='EncoderDecoder',
    pretrained=
    'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_small_p16_384_20220308-410f6037.pth',
    backbone=dict(
        type='VisionTransformer',
        img_size=(512, 512),
        patch_size=16,
        in_channels=3,
        embed_dims=384,
        num_layers=12,
        num_heads=6,
        drop_path_rate=0.1,
        attn_drop_rate=0.0,
        drop_rate=0.0,
        final_norm=True,
        norm_cfg=dict(type='LN', eps=1e-06, requires_grad=True),
        with_cls_token=True,
        interpolate_mode='bicubic'),
    decode_head=dict(
        type='SegmenterMaskTransformerHead',
        in_channels=384,
        channels=384,
        num_classes=6,
        num_layers=2,
        num_heads=6,
        embed_dims=384,
        dropout_ratio=0.0,
        loss_decode=dict(
            type='CrossEntropyLoss',
            use_sigmoid=False,
            loss_weight=1.0,
            avg_non_ignore=True,
            class_weight=[0.624, 0.966, 0.772, 0.821, 0.872, 0.945]),
        ignore_index=255),
    test_cfg=dict(mode='slide', crop_size=(512, 512), stride=(480, 480)))
log_config = dict(
    interval=50, hooks=[dict(type='TextLoggerHook', by_epoch=False)])
dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]
cudnn_benchmark = True
optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0)
optimizer_config = dict()
lr_config = dict(policy='poly', power=0.9, min_lr=0.0001, by_epoch=False)
runner = dict(type='IterBasedRunner', max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=2000)
evaluation = dict(interval=500, metric='mIoU', pre_eval=True, save_best='mIoU')
dataset_type = 'CustomDataset'
data_root = 'e3_training_set_v17'
classes = ['ground', 'tree', 'sky', 'trail', 'rock', 'bush']
class_weights = [0.624, 0.966, 0.772, 0.821, 0.872, 0.945]
test_palette = [(255, 218, 165), (137, 189, 134), (195, 243, 255),
                (255, 255, 186), (199, 206, 234), (226, 240, 203)]
img_norm_cfg = dict(
    mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], to_rgb=True)
crop_size = (512, 512)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=False),
    dict(type='Resize', img_scale=(1440, 900), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=(512, 512), cat_max_ratio=0.5),
    dict(type='RandomFlip', prob=0.5),
    dict(
        type='Normalize',
        mean=[127.5, 127.5, 127.5],
        std=[127.5, 127.5, 127.5],
        to_rgb=True),
    dict(type='Pad', size=(512, 512), pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg'])
]
#test_pipeline = [
#    dict(type='LoadImageFromFile'),
#    dict(
#        type='MultiScaleFlipAug',
#        img_scale=(1440, 900),
#        flip=False,
#        transforms=[
#            dict(
#                type='Normalize',
#                mean=[127.5, 127.5, 127.5],
#                std=[127.5, 127.5, 127.5],
#                to_rgb=True),
#            dict(type='ImageToTensor', keys=['img']),
#            dict(type='Collect', keys=['img'])
#        ])
#]
data = dict(
    samples_per_gpu=4,
    workers_per_gpu=4,
    train=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        'e3_training_set_v17',
        img_dir='images/train',
        ann_dir='labels/train',
        pipeline=[
            dict(type='LoadImageFromFile'),
	     dict(
		 type='MultiScaleFlipAug',
		 img_scale=(1440, 900),
		 flip=False,
		 transforms=[
		     dict(type='Resize', keep_ratio=True),
		     dict(
		         type='Normalize',
		         mean=[127.5, 127.5, 127.5],
		         std=[127.5, 127.5, 127.5],
		         to_rgb=True),
		     dict(type='ImageToTensor', keys=['img']),
		     dict(type='Collect', keys=['img'])
		 ]
	    )
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush']),
    val=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        'e3_training_set_v17',
        img_dir='images/val',
        ann_dir='labels/val',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1440, 900),
                flip=False,
                transforms=[
                    dict(
                        type='Normalize',
                        mean=[127.5, 127.5, 127.5],
                        std=[127.5, 127.5, 127.5],
                        to_rgb=True),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush']),
    test=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        'e3_training_set_v17',
        img_dir='images/val',
        ann_dir='labels/val',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1440, 900),
                flip=False,
                transforms=[
                    dict(
                        type='Normalize',
                        mean=[127.5, 127.5, 127.5],
                        std=[127.5, 127.5, 127.5],
                        to_rgb=True),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush'],
        palette=[(255, 218, 165), (137, 189, 134), (195, 243, 255),
                 (255, 255, 186), (199, 206, 234), (226, 240, 203)]))
work_dir = './work_dirs/e3_segmenter'
gpu_ids = [0]
auto_resume = False

[irexyc]

Not data.train.pipeline but data.test.pipeline.

Below is my env and config. I can convert it successful.

mmseg:	0.24.1
MMDeploy: 0.6.0+5195ff9
MMCV: 1.6.0
PyTorch: 1.8.0+cu111
TorchVision: 0.9.0+cu111

checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_small_p16_384_20220308-410f6037.pth'
backbone_norm_cfg = dict(type='LN', eps=1e-06, requires_grad=True)
model = dict(
    type='EncoderDecoder',
    pretrained=
    'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_small_p16_384_20220308-410f6037.pth',
    backbone=dict(
        type='VisionTransformer',
        img_size=(512, 512),
        patch_size=16,
        in_channels=3,
        embed_dims=384,
        num_layers=12,
        num_heads=6,
        drop_path_rate=0.1,
        attn_drop_rate=0.0,
        drop_rate=0.0,
        final_norm=True,
        norm_cfg=dict(type='LN', eps=1e-06, requires_grad=True),
        with_cls_token=True,
        interpolate_mode='bicubic'),
    decode_head=dict(
        type='SegmenterMaskTransformerHead',
        in_channels=384,
        channels=384,
        num_classes=6,
        num_layers=2,
        num_heads=6,
        embed_dims=384,
        dropout_ratio=0.0,
        loss_decode=dict(
            type='CrossEntropyLoss',
            use_sigmoid=False,
            loss_weight=1.0,
            avg_non_ignore=True,
            class_weight=[0.624, 0.966, 0.772, 0.821, 0.872, 0.945]),
        ignore_index=255),
    test_cfg=dict(mode='slide', crop_size=(512, 512), stride=(480, 480)))
log_config = dict(
    interval=50, hooks=[dict(type='TextLoggerHook', by_epoch=False)])
dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]
cudnn_benchmark = True
optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0)
optimizer_config = dict()
lr_config = dict(policy='poly', power=0.9, min_lr=0.0001, by_epoch=False)
runner = dict(type='IterBasedRunner', max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=2000)
evaluation = dict(interval=500, metric='mIoU', pre_eval=True, save_best='mIoU')
dataset_type = 'CustomDataset'
data_root = '/training_set_v17'
classes = ['ground', 'tree', 'sky', 'trail', 'rock', 'bush']
class_weights = [0.624, 0.966, 0.772, 0.821, 0.872, 0.945]
test_palette = [(255, 218, 165), (137, 189, 134), (195, 243, 255),
                (255, 255, 186), (199, 206, 234), (226, 240, 203)]
img_norm_cfg = dict(
    mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], to_rgb=True)
crop_size = (512, 512)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=False),
    dict(type='Resize', img_scale=(1440, 900), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=(512, 512), cat_max_ratio=0.5),
    dict(type='RandomFlip', prob=0.5),
    dict(
        type='Normalize',
        mean=[127.5, 127.5, 127.5],
        std=[127.5, 127.5, 127.5],
        to_rgb=True),
    dict(type='Pad', size=(512, 512), pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg'])
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(1440, 900),
        flip=False,
        transforms=[
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_rgb=True),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img'])
        ])
]
data = dict(
    samples_per_gpu=4,
    workers_per_gpu=4,
    train=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        '/training_set_v17',
        img_dir='images/train',
        ann_dir='labels/train',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(type='LoadAnnotations', reduce_zero_label=False),
            dict(type='Resize', img_scale=(1440, 900), ratio_range=(0.5, 2.0)),
            dict(type='RandomCrop', crop_size=(512, 512), cat_max_ratio=0.5),
            dict(type='RandomFlip', prob=0.5),
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_rgb=True),
            dict(type='Pad', size=(512, 512), pad_val=0, seg_pad_val=255),
            dict(type='DefaultFormatBundle'),
            dict(type='Collect', keys=['img', 'gt_semantic_seg'])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush']),
    val=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        '/training_set_v17',
        img_dir='images/val',
        ann_dir='labels/val',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1440, 900),
                flip=False,
                transforms=[
                    dict(
                        type='Normalize',
                        mean=[127.5, 127.5, 127.5],
                        std=[127.5, 127.5, 127.5],
                        to_rgb=True),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush']),
    test=dict(
        type='CustomDataset',
        reduce_zero_label=False,
        data_root=
        '/training_set_v17',
        img_dir='images/val',
        ann_dir='labels/val',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1440, 900),
                flip=False,
                transforms=[
                    dict(type='Resize', keep_ratio=True),
                    dict(
                        type='Normalize',
                        mean=[127.5, 127.5, 127.5],
                        std=[127.5, 127.5, 127.5],
                        to_rgb=True),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ],
        classes=['ground', 'tree', 'sky', 'trail', 'rock', 'bush'],
        palette=[(255, 218, 165), (137, 189, 134), (195, 243, 255),
                 (255, 255, 186), (199, 206, 234), (226, 240, 203)]))
work_dir = './work_dirs/segmenter'
gpu_ids = [0]
auto_resume = False

[sammilei]

@irexyc Thanks a lot for providing your replies!
Note My model is fine-tuned on the mask vit. I think my issue is with the mask vit config and it fails with tensors being in GPU and CPU. Could you try case 3? See my test below:

1. [worked] linear vit: mmseg config + mmseg pth: as in [Bug] IShuffleLayer applied to shape tensor must have 0 or 1 reshape dimensions: dimensions were [-1,2]) #2098
2. linear vit: your config + mmseg pth: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!.
3. mask vit: mmseg config + mmseg pth: RuntimeError: Expected all tensors to be on the same device...
4. mask vit: your config + mmseg pth: RuntimeError: Expected all tensors to be on the same device...
5. mask vit: your config + my pth: RuntimeError: Expected all tensors to be on the same device...

My env:

mmseg:	0.24.1
MMDeploy: 0.6.0
MMCV: 1.6.0
PyTorch: 1.12.0a0+2c916ef.nv22.3
TorchVision: 0.13.0
torch.cuda.is_available(): True

Not sure if I missed any setting so the tensors would not go to CPU.