PyTorch：灵活加载指定层参数

Pytorch加载指定层的参数
在PyTorch中，我们经常需要加载预训练模型的参数以用于迁移学习或其他任务。 有时，我们可能只对模型的一部分感兴趣，例如只加载模型的某些层。 下面是如何做到这一点的步骤。
在PyTorch中，序列化模型的状态字典可以通过torch.load()函数加载。 默认情况下，此函数将加载整个模型，包括所有层。 如果你只想要加载一部分模型，你可以通过以下方式做到：

import torch
import torchvision.models as models
# 加载预训练模型
model = models.resnet50(pretrained=True)
# 如果你想要加载模型的特定层，你可以创建一个新的模型，然后从预训练的模型复制所需的层。
new_model = models.resnet50(pretrained=False)  # 创建一个新的模型
# 现在我们只复制我们想要的层。 我们可以通过遍历模型的参数并检查其名字来实现这一点。
for name, param in model.named_parameters():
if 'layer3' in name or 'layer4' in name:  # 例如，我们想要复制'layer3'和'layer4'的所有参数
new_param = param.data.clone()  # 复制参数
new_model.state_dict()[name] = new_param  # 将新的参数放入我们的新模型中
# 现在，new_model只有我们感兴趣的层的参数。

上面的代码展示了如何从预训练的模型中提取并加载指定层的参数。在实际应用中，你可能需要根据自己的需求来调整代码，例如根据不同的层名或结构特征进行筛选。
注意：在复制参数时，请确保你复制的是参数的数据（即param.data），而不是参数本身（即param）。这是因为PyTorch的参数对象是不可变的，你不能直接修改它的值。但是，你可以修改参数的数据。
此外，如果你正在处理一个很大的模型，你可能需要使用一些优化方法来减少内存使用，例如使用torch.utils.checkpoint.checkpoint或torch.utils.checkpoint.checkpoints_io。这些方法可以减少在复制大量数据时使用的内存，特别是当你只对模型的最后几层感兴趣时。