|
|
| import argparse
|
| from collections import OrderedDict
|
| import torch
|
|
|
|
|
| def convert(src: str, dst: str):
|
| checkpoint = torch.load(src)
|
| has_model = "model" in checkpoint.keys()
|
| checkpoint = checkpoint["model"] if has_model else checkpoint
|
| if "state_dict" in checkpoint.keys():
|
| checkpoint = checkpoint["state_dict"]
|
| out_cp = OrderedDict()
|
| for k, v in checkpoint.items():
|
| out_cp[".".join(["backbone", k])] = v
|
| out_cp = {"model": out_cp} if has_model else out_cp
|
| torch.save(out_cp, dst)
|
|
|
|
|
| if __name__ == "__main__":
|
| parser = argparse.ArgumentParser()
|
| parser.add_argument(
|
| "--src", "-s", type=str, required=True, help="Path to src weights.pth"
|
| )
|
| parser.add_argument(
|
| "--dst", "-d", type=str, required=True, help="Path to dst weights.pth"
|
| )
|
| args = parser.parse_args()
|
| convert(
|
| args.src,
|
| args.dst,
|
| )
|
|
|
