使用自己预训练好的模型作为网络的初始化参数是深度学习中常见的做法,尤其是在迁移学习场景中。这可以通过加载预训练模型的权重到新模型中来实现,前提是新模型的架构能够与预训练模型的权重相匹配。以下是一般步骤和示例代码,展示如何在PyTorch中实现这一过程。
步骤1:保存预训练模型的权重
假设你已经训练了一个模型,并且想要保存它的权重,可以使用以下代码:
torch.save(model.state_dict(), 'model_weights.pth')
这里,model
是你的预训练模型,model_weights.pth
是保存模型权重的文件。
步骤2:加载权重到新模型
当你有了预训练模型的权重后,可以将这些权重加载到新的模型中,前提是新模型的架构必须与预训练模型兼容。如果新模型与预训练模型完全相同,可以直接加载权重。如果只是部分相同,你可能需要在加载权重之前进行一些调整。