lyuwenyu / RT-DETR

[CVPR 2024] Official RT-DETR (RTDETR paddle pytorch), Real-Time DEtection TRansformer, DETRs Beat YOLOs on Real-time Object Detection. 🔥 🔥 🔥
Apache License 2.0
1.87k stars 192 forks source link

如何将paddle版本的pretrain_weights转化为pytorch版本的? #352

Open hzy-del opened 2 weeks ago

lyuwenyu commented 2 weeks ago

https://github.com/lyuwenyu/cvperception/blob/main/src/cvperception/zoo/rtdetr/conver_params.py

"""Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""

def main(args) -> None:
    import torch 
    from cvperception.core import YAMLConfig

    model = YAMLConfig(args.config).model

    if args.version == 1:
        state = model.state_dict()
        keys = [k for k in state.keys() if 'num_batches_tracked' not in k]

    elif args.version == 2:
        state = model.state_dict()
        ignore_keys = ['anchors', 'valid_mask', 'num_points_scale']
        keys = [k for k in state.keys() if 'num_batches_tracked' not in k]
        keys = [k for k in keys if not any([x in k for x in ignore_keys])]

    import paddle
    p_state = paddle.load(args.pdparams)
    pkeys = list(p_state.keys())

    assert len(keys) == len(pkeys), f'{len(keys)}, {len(pkeys)}'

    new_state = {}
    for i, k in enumerate(keys):    
        pp = p_state[pkeys[i]]
        pp = torch.tensor(pp.numpy())

        if 'denoising_class_embed' in k:
            new_state[k] = torch.concat([pp, torch.zeros(1, pp.shape[-1])], dim=0)
            continue

        tp = state[k]
        if len(tp.shape) == 2:
            new_state[k] = pp.T
        elif len(tp.shape) == 1:
            new_state[k] = pp
        else:
            assert tp.shape == pp.shape, f'{k}, {pp.shape}, {tp.shape}'
            new_state[k] = pp

    assert len(new_state) == len(p_state), ''

    # checkpoint = {'ema': {'module': new_state, }}
    # torch.save(checkpoint, args.output_file)

    model.load_state_dict(new_state, strict=False)

    checkpoint = {'ema': {'module': model.state_dict(), }}
    torch.save(checkpoint, args.output_file)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, )
    parser.add_argument('-p', '--pdparams', type=str, )
    parser.add_argument('-o', '--output_file', type=str, )
    parser.add_argument('-v', '--version', type=int, default=1)

    args = parser.parse_args()
    main(args)
hzy-del commented 2 weeks ago
from cvperception.core import YAMLConfig

非常感谢!!我还想问一下YAMLConfig文件的配置是怎样的呢?

lyuwenyu commented 2 weeks ago

这个repo从这里导入 https://github.com/lyuwenyu/RT-DETR/blob/main/rtdetr_pytorch/src/core/yaml_config.py#L14