最佳想把在torchversion库加载的预训练权重基础上训练过的模型权重加载到Detectron2框架的任务中进行实验。
发现两者必须经过一个转换,实验过后有效的转换方式如下所示,
#!/usr/bin/env python
import pickle as pkl
import sysimport torchif __name__ == "__main__":input = "./model_best.pth.tar"obj = torch.load(input, map_location="cpu")obj = obj["state_dict"]newmodel = {}for k, v in obj.items():old_k = kif "layer" not in k:k = "stem." + kfor t in [1, 2, 3, 4]:k = k.replace("layer{}".format(t), "res{}".format(t + 1))for t in [1, 2, 3]:k = k.replace("bn{}".format(t), "conv{}.norm".format(t))k = k.replace("downsample.0", "shortcut")k = k.replace("downsample.1", "shortcut.norm")print(old_k,v.shape, "->", k)newmodel[k] = vres = {"model": newmodel}torch.save(res,"./output.pth")
其中我发现当参考网上大多数的解答,使用下面的方式保存的权重,转后后载入时会报错:RuntimeError: Invalid magic number; corrupt file?
换成上面的torch.save就能解决
res = {"model": newmodel}with open("./output.pth", "wb") as f:pkl.dump(res, f)