PyTorch加载预训练模型(pretrained)

最新推荐文章于 2024-01-25 21:10:52 发布

原创最新推荐文章于 2024-01-25 21:10:52 发布 · 1.1w 阅读

32 ·

CC 4.0 BY-SA版权

文章标签：

#加载预训练模型 #PyTorch

项目实现专栏收录该内容

3 篇文章

订阅专栏

本文详细介绍了如何使用PyTorch加载预训练的ResNet50模型，并通过微调模型参数，使其适用于特定任务。包括加载模型、读取预训练参数、更新模型参数等关键步骤。

部署运行你感兴趣的模型镜像

使用预训练模型的代码如下：

# 加载预训练模型
    resNet50 = models.resnet50(pretrained=True)
    ResNet50 = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=2)

    # 读取参数
    pretrained_dict = resNet50.state_dict()
    model_dict = ResNet50.state_dict()

    # 将pretained_dict里不属于model_dict的键剔除掉
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

    # 更新现有的model_dict
    model_dict.update(pretrained_dict)

    # 加载真正需要的state_dict
    ResNet50.load_state_dict(model_dict)

您可能感兴趣的与本文相关的镜像