pytorch：加载预训练模型中的部分参数，并固定该部分参数

转载已于 2022-02-23 17:14:36 修改 · 584 阅读

1 ·

本内容遵循CC 4.0 BY-SA版权协议

原文链接：https://www.jianshu.com/p/d67d62982a24

收录于

当前文章被以下社区和专栏收录：

于 2022-02-23 17:13:49 首次发布

这篇博客介绍了如何在PyTorch中使用DenseNet121模型，仅加载和固定预训练模型的部分参数。首先，定义并初始化模型，然后将模型转移到CUDA设备上。接着，通过更新模型状态字典并加载预训练权重，实现了对部分参数的加载。最后，使用DataParallel进行并行计算。

pytorch：加载预训练模型中的部分参数，并固定该部分参数

https://www.jianshu.com/p/d67d62982a24

		initial_cnn = models.densenet121(pretrained=False)
        self.cnn = torch.nn.Sequential(*(list(initial_cnn.children())[:-1]))
        device = torch.device('cuda')
        cnn = self.cnn.to(device)
        cnn = nn.DataParallel(cnn)

        model_dict = cnn.state_dict()

        pretrained_dict = torch.load('epoch_11.pth')

        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

        model_dict.update(pretrained_dict)

        cnn.load_state_dict(model_dict)