pytorch加载部分权重的方法

刀么克瑟拉莫

于 2023-06-13 14:38:13 发布

阅读量272

点赞数

分类专栏： pytorch 文章标签： pytorch

本文链接：https://blog.youkuaiyun.com/random_repick/article/details/131187882

版权

pytorch 专栏收录该内容

17 篇文章

订阅专栏

该代码段展示了如何在PyTorch中将预训练的模型权重加载到神经网络结构中。首先定义网络结构，然后根据CUDA是否可用决定模型运行的设备。接着从weights.pth文件中加载状态字典，并逐层地将权重加载到网络的子模块中，严格匹配参数。

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

import torch

net = yourNet()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
state=torch.load('weights.pth', map_location=device)
for name, child in net.named_children():
	if name in state.keys():
		child.load_state_dict(state[name], strict=True)