本文共 749 字，大约阅读时间需要 2 分钟。

加载预训练模型并初始化新模型权重

在PyTorch中加载预训练模型并初始化新模型的过程中，我们可以参考以下步骤：

加载预训练模型

首先，从指定路径加载预训练模型的状态字典：

saved = torch.load(cfg.Transfer, map_location=device)
old_state_dict = saved['state_dict']

初始化新模型

创建一个空的字典用于存储新模型的权重：

new_state_dict = {}

处理模型层

遍历当前模型的状态字典，并将旧模型中与新模型层形状匹配的权重复制到新模型中：

all_layer = len(model.state_dict())
num = 0
for key in model.state_dict():
    if key in old_state_dict and old_state_dict[key].shape == model.state_dict()[key].shape:
        new_state_dict[key] = old_state_dict[key]
        num += 1
    else:
        new_state_dict[key] = model.state_dict()[key]

加载新模型部分权重

使用load_state_dict方法加载新模型的部分权重：

model.load_state_dict(new_state_dict)

输出结果

打印加载的层数信息：

print("从预训练模型中加载了{}/{}层".format(num, all_layer))

这个方法可以帮助开发者在保持模型灵活性的同时，充分利用预训练模型的优势。

转载地址：http://clxfk.baihongyu.com/