Spaces:
Runtime error
Runtime error
File size: 2,178 Bytes
583456e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
# Copyright (c) Facebook, Inc. and its affiliates.
# Copyright (c) Meta Platforms, Inc. All Rights Reserved
import pickle as pkl
import sys
import torch
"""
Usage:
# download pretrained swin model:
wget https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth
# run the conversion
./convert-pretrained-model-to-d2.py swin_tiny_patch4_window7_224.pth swin_tiny_patch4_window7_224.pkl
# Then, use swin_tiny_patch4_window7_224.pkl with the following changes in config:
MODEL:
WEIGHTS: "/path/to/swin_tiny_patch4_window7_224.pkl"
INPUT:
FORMAT: "RGB"
"""
def transform(path):
model = torch.load(path, map_location="cpu")
print(f"loading {path}......")
state_dict = model["model"]
state_dict = {
k.replace("visual_model.", ""): v
for k, v in state_dict.items()
if k.startswith("visual_model")
}
source_keys = [k for k in state_dict.keys() if "relative_coords" in k]
for k in source_keys:
state_dict[
k.replace("relative_coords", "relative_position_index")
] = state_dict[k]
del state_dict[k]
source_keys = [k for k in state_dict.keys() if "atten_mask_matrix" in k]
for k in source_keys:
state_dict[k.replace("atten_mask_matrix", "attn_mask")] = state_dict[k]
del state_dict[k]
source_keys = [k for k in state_dict.keys() if "rel_pos_embed_table" in k]
for k in source_keys:
state_dict[
k.replace("rel_pos_embed_table", "relative_position_bias_table")
] = state_dict[k]
del state_dict[k]
source_keys = [k for k in state_dict.keys() if "channel_reduction" in k]
for k in source_keys:
state_dict[k.replace("channel_reduction", "reduction")] = state_dict[k]
del state_dict[k]
return {
k if k.startswith("backbone.") else "backbone." + k: v
for k, v in state_dict.items()
}
if __name__ == "__main__":
input = sys.argv[1]
res = {
"model": transform(input),
"__author__": "third_party",
"matching_heuristics": True,
}
with open(sys.argv[2], "wb") as f:
pkl.dump(res, f)
|