首页 > 其他 > 详细

torch.load加载权重时报错 Magic Number Error

时间:2020-07-19 11:31:25      阅读:88      评论:0      收藏:0      [点我收藏+]

有时候使用 torch.load 加载比较古老的权重文件时可能报错 Magic Number Error,这有可能是因为该文件使用 pickle 存储并且编码使用了 latin1,此时可以这样加载:

def resnet50(weights_path=None, **kwargs):
    """Constructs a ResNet-50 model.
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    if weights_path:
        import pickle
        with open(weights_path, rb) as f:
            obj = f.read()
        weights = {key: weight_dict for key, weight_dict in pickle.loads(obj, encoding=latin1).items()}
        model.load_state_dict(weights)
    return model

参考:

https://github.com/cydonia999/VGGFace2-pytorch/issues/2#issuecomment-463571389

torch.load加载权重时报错 Magic Number Error

原文:https://www.cnblogs.com/luruiyuan/p/13338219.html

(0)
(0)
   
举报
评论 一句话评论(0
关于我们 - 联系我们 - 留言反馈 - 联系我们:wmxa8@hotmail.com
© 2014 bubuko.com 版权所有
打开技术之扣,分享程序人生!