关于Pytorch加载预训练模型

vectorindex

于 2021-04-11 14:55:57 发布

阅读量1k

点赞数

CC 4.0 BY-SA版权

分类专栏：机器学习

本文链接：https://blog.youkuaiyun.com/vectorindex/article/details/115599276

本文介绍了在Pytorch中加载预训练模型时需要注意的问题，包括自定义网络类的权重输出顺序、模型与参数匹配的重要性，以及如何处理变量命名不一致的情况。确保.pth文件中的字典顺序、权重形状和变量名与自建网络完全一致才能成功加载。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1.关于Pytorch加载预训练模型

自定义网络类的字典是OrderedDict字典，为顺序的，即是按照代码中声明层的先后顺序输出的。

    def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
                               kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU()
        #----------------------18，34是由2个一样的卷积核组成---------------
        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
                               kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)