Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| import argparse | |
| import pickle | |
| import torch | |
| """ | |
| Usage: | |
| cd DETIC_ROOT/models/ | |
| wget https://miil-public-eu.oss-eu-central-1.aliyuncs.com/model-zoo/ImageNet_21K_P/models/resnet50_miil_21k.pth | |
| python ../tools/convert-thirdparty-pretrained-model-to-d2.py --path resnet50_miil_21k.pth | |
| wget https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth | |
| python ../tools/convert-thirdparty-pretrained-model-to-d2.py --path swin_base_patch4_window7_224_22k.pth | |
| """ | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--path', default='') | |
| args = parser.parse_args() | |
| print('Loading', args.path) | |
| model = torch.load(args.path, map_location="cpu") | |
| # import pdb; pdb.set_trace() | |
| if 'model' in model: | |
| model = model['model'] | |
| if 'state_dict' in model: | |
| model = model['state_dict'] | |
| ret = { | |
| "model": model, | |
| "__author__": "third_party", | |
| "matching_heuristics": True | |
| } | |
| out_path = args.path.replace('.pth', '.pkl') | |
| print('Saving to', out_path) | |
| pickle.dump(ret, open(out_path, "wb")) | |