关于Pytorch加载预训练模型

最新推荐文章于 2024-07-26 20:52:07 发布

原创

最新推荐文章于 2024-07-26 20:52:07 发布 · 1k 阅读

4 ·

CC 4.0 BY-SA版权

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

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)