How can I make simcc Head + HRNet config?

[seon-creator]

Hello, When I checked model zoo in Body 2d keypoints, I saw a Simcc paper.
In that paper, the result was Simcc Head,HRNet method was better than Topdown heatmap HRNet method.
So, I want to train Simcc+HRNet, but in mmpose, there are not HRNet config in Simcc directory.

So, I made custom config with reference of mmpose/configs/body_2d_keypoint/simcc/coco/simcc_res50_8xb32-140e_coco-384x288.py.
I copy the topdown heatmap HRNet config, and changed codec and head to Simcc.

This is the custom config code of my Simcc HRNet.

_base_ = ['../../../_base_/default_runtime.py']

# runtime
train_cfg = dict(max_epochs=210, val_interval=10)

# optimizer
optim_wrapper = dict(optimizer=dict(
    type='Adam',
    lr=1e-3,
))

# learning policy
param_scheduler = [
    dict(
        type='LinearLR', begin=0, end=500, start_factor=0.001,
        by_epoch=False),  # warm-up
    dict(
        type='MultiStepLR',
        begin=0,
        end=300,    #edit from 210
        milestones=[170, 200],
        gamma=0.1,
        by_epoch=True)
]

# automatically scaling LR based on the actual training batch size
auto_scale_lr = dict(base_batch_size=512)

# hooks
default_hooks = dict(checkpoint=dict(save_best='coco/AP', rule='greater'))

# Change
# codec settings
codec = dict(
    type='SimCCLabel',
    input_size=(256, 256),
    sigma=6.0, 
    simcc_split_ratio=2.0,
    use_dark = True # edit
    )

# model settings
model = dict(
    type='TopdownPoseEstimator',
    # data preprocessor
    data_preprocessor=dict(
        type='PoseDataPreprocessor',
        mean=[123.675, 116.28, 103.53],
        std=[58.395, 57.12, 57.375],
        bgr_to_rgb=True),
    # backbone
    backbone=dict(
        type='HRNet',
        in_channels=3,
        extra=dict(
            stage1=dict(
                num_modules=1,
                num_branches=1,
                block='BOTTLENECK',
                num_blocks=(4, ),
                num_channels=(64, )),
            stage2=dict(
                num_modules=1,
                num_branches=2,
                block='BASIC',
                num_blocks=(4, 4),
                num_channels=(32, 64)),
            stage3=dict(
                num_modules=4,
                num_branches=3,
                block='BASIC',
                num_blocks=(4, 4, 4),
                num_channels=(32, 64, 128)),
            stage4=dict(
                num_modules=3,
                num_branches=4,
                block='BASIC',
                num_blocks=(4, 4, 4, 4),
                num_channels=(32, 64, 128, 256))),
        # Using pretrained model to get good performance with small dataset
        # init_cfg=dict(
        #     type='Pretrained',
        #     # If you want to load a checkpoint to your backbone, you should specify the prefix
        #     # prefix='backbone.',
        #     checkpoint='/data/home/seondeok/Project/acupoint/mmpose/configs/body_2d_keypoint/pretrain/simcc/pose_hrnet_w32_256x192.pth'
        #     ),
    ),
    # Change
    head=dict(
        type='SimCCHead',
        in_channels=32,
        out_channels=5, # edit
        input_size=codec['input_size'],
        in_featuremap_size=tuple([s // 32 for s in codec['input_size']]),
        simcc_split_ratio=codec['simcc_split_ratio'],
        loss=dict(type='KLDiscretLoss', use_target_weight=True),
        decoder=codec),
    test_cfg=dict(flip_test=True))

# base dataset settings
dataset_type = 'CocoArm'
data_mode = 'topdown'
data_root = '/data/home/seondeok/Project/acupoint/coco_dataset/arm/dataset3/Train760/'
test_root = '/data/home/seondeok/Project/acupoint/coco_dataset/arm/dataset3/Test180/'
annotation_root = '/data/home/seondeok/Project/acupoint/coco_dataset/arm/dataset3/Train760.json'
annotation_root_val = '/data/home/seondeok/Project/acupoint/coco_dataset/arm/dataset3/Test180.json'

# pipelines
train_pipeline = [
    dict(type='LoadImage'),
    dict(type='GetBBoxCenterScale'),
    dict(type='RandomFlip', direction='diagonal'),
    dict(type='RandomHalfBody'),
    dict(type='RandomBBoxTransform'),
    dict(type='TopdownAffine', input_size=codec['input_size']),
    dict(type='GenerateTarget', encoder=codec),
    dict(type='PackPoseInputs')
]
val_pipeline = [
    dict(type='LoadImage'),
    dict(type='GetBBoxCenterScale'),
    dict(type='TopdownAffine', input_size=codec['input_size']),
    dict(type='PackPoseInputs')
]

# data loaders
train_dataloader = dict(
    batch_size=16,
    num_workers=2,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        data_mode=data_mode,
        ann_file=annotation_root,
        data_prefix=dict(img=data_root),
        pipeline=train_pipeline,
    ))
val_dataloader = dict(
    batch_size=16,
    num_workers=2,
    persistent_workers=True,
    drop_last=False,
    sampler=dict(type='DefaultSampler', shuffle=False, round_up=False),
    dataset=dict(
        type=dataset_type,
        data_root=test_root,
        data_mode=data_mode,
        ann_file=annotation_root_val,
        # bbox_file='data/coco/person_detection_results/'
        # 'COCO_val2017_detections_AP_H_56_person.json',
        data_prefix=dict(img=test_root),
        test_mode=True,
        pipeline=val_pipeline,
    ))
test_dataloader = val_dataloader

# evaluators
val_evaluator = dict(
    type='CocoMetric',
    ann_file=annotation_root_val)
test_evaluator = val_evaluator

I marked that I changed.

And I train the model, but in this step, This error accur.
How to solve the problem?

Traceback (most recent call last):
  File "tools/train.py", line 161, in <module>
    main()
  File "tools/train.py", line 157, in main
    runner.train()
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/runner.py", line 1745, in train
    model = self.train_loop.run()  # type: ignore
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/loops.py", line 96, in run
    self.run_epoch()
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/loops.py", line 112, in run_epoch
    self.run_iter(idx, data_batch)
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/loops.py", line 128, in run_iter
    outputs = self.runner.model.train_step(
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/mmengine/model/base_model/base_model.py", line 114, in train_step
    losses = self._run_forward(data, mode='loss')  # type: ignore
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/mmengine/model/base_model/base_model.py", line 340, in _run_forward
    results = self(**data, mode=mode)
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/home/seondeok/Project/acupoint/mmpose/mmpose/models/pose_estimators/base.py", line 136, in forward
    return self.loss(inputs, data_samples)
  File "/data/home/seondeok/Project/acupoint/mmpose/mmpose/models/pose_estimators/topdown.py", line 74, in loss
    self.head.loss(feats, data_samples, train_cfg=self.train_cfg))
  File "/data/home/seondeok/Project/acupoint/mmpose/mmpose/models/heads/coord_cls_heads/simcc_head.py", line 323, in loss
    pred_x, pred_y = self.forward(feats)
  File "/data/home/seondeok/Project/acupoint/mmpose/mmpose/models/heads/coord_cls_heads/simcc_head.py", line 223, in forward
    pred_x = self.mlp_head_x(x)
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/home/seondeok/.conda/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (80x262144 and 4096x512)