2020-12-10 16:57:52 +00:00
|
|
|
import torch
|
|
|
|
|
|
2020-12-18 16:24:31 +00:00
|
|
|
from models.spinenet_arch import SpineNet
|
2020-12-10 16:57:52 +00:00
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
2021-01-06 21:52:17 +00:00
|
|
|
pretrained_path = '../../../experiments/resnet_byol_diffframe_115k.pth'
|
2021-01-06 03:14:22 +00:00
|
|
|
output_path = '../../experiments/resnet_byol_diffframe_115k_.pth'
|
2020-12-10 16:57:52 +00:00
|
|
|
|
|
|
|
|
wrap_key = 'online_encoder.net.'
|
|
|
|
|
sd = torch.load(pretrained_path)
|
|
|
|
|
sdo = {}
|
|
|
|
|
for k,v in sd.items():
|
|
|
|
|
if wrap_key in k:
|
|
|
|
|
sdo[k.replace(wrap_key, '')] = v
|
|
|
|
|
|
2021-01-03 23:37:18 +00:00
|
|
|
#model = SpineNet('49', in_channels=3, use_input_norm=True).to('cuda')
|
|
|
|
|
#model.load_state_dict(sdo, strict=True)
|
2020-12-10 16:57:52 +00:00
|
|
|
|
|
|
|
|
print("Validation succeeded, dumping state dict to output path.")
|
|
|
|
|
torch.save(sdo, output_path)
|
