Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 681 Bytes
6dfcb0f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
import torch
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, help='The path to the checkpoint.')
parser.add_argument('--output', type=str, default=None, help='the output path of the checkpoint')
args = parser.parse_args()
state_dict = torch.load(args.input, map_location='cpu')['model']
new_state_dict = {}
for k, v in state_dict.items():
if 'encoder' in k and not 'decoder' in k:
new_k = 'backbone.net.model.' + k
new_state_dict[new_k] = v
output_path = args.input.replace('.pth', '-encoder.pth') if args.output is None else args.output
torch.save(new_state_dict, output_path)
print('Save model to', output_path) |