裸问题陈述:
我训练了一个模型 A，它由特征提取器 FE 和分类头 ACH 组成。
我想训练一个模型 B，它使用 A 的特征提取器 FE 并重新训练它自己的分类头 BCH。
到目前为止，这很容易。现在我不想保存整个模型B，因为它的FE部分已经保存在模型A中。我只想转储BCH，并且在推理期间
- 加载模型 A - 进行预测
- 加载B的分类头BCH。
- 将分类头 ACH 与 BCH 交换
- 使用此交换状态运行预测。
阅读pyTorches documentation它只讨论保存整个模型。我怎样才能实现这个目标？
问题陈述结束
有关问题动机的更多详细信息:
我有一个想要分类的图像数据集，这些图像可以有多个类别。例如，同一图像可以具有“陆地车辆”类别(父类(super class)别)和“汽车”类别(类别)或“卡车”类别。另一张图像可能具有“飞行器”类别，并且可以是“直升机”或“飞机”。
由于图像和大多数特征应该是相同的，我希望为父类(super class)别训练一个分类器，然后卡住它的特征提取器，并使用预训练的特征提取器为类别转移学习相同的模型.
由于特征提取主干的权重是相同的，所以我只想保存类别模型的分类头的权重，从而节省一些宝贵的计算资源。
最佳答案
一般来说，通常只需要访问模型的主干，以便将其重用于其他目的。您有多种方法可以执行此操作。但最重要的是，要记住保存模型检查点并稍后加载意味着保存权重和偏差并能够将它们正确加载到相应的层，您首先需要从模型中知道您想要保存的部分。
当你获得模型的状态时，你将获得一个字典。 键将是层名称，值将是权重和偏差。让我们看一个带有 efficientnet 的示例分类器如何仅保存模型的主干。基本上，一个高效网络，如您的示例中所示，是一个主干和一个作为头部的全连接层，如果您只想要主干，那么您需要每个层，除了稍后要微调的头部。
import torch
import torch.nn as nn
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_name("efficientnet-b0")
print(model)
它将打印模型层和一些功能，基本的东西。
EfficientNet(
(_conv_stem): Conv2dStaticSamePadding(
3, 32, kernel_size=(3, 3), stride=(2, 2), bias=False
(static_padding): ZeroPad2d(padding=(0, 1, 0, 1), value=0.0)
)
(_bn0): BatchNorm2d(32, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)
(_blocks): ModuleList(
(0): MBConvBlock(
(_depthwise_conv): Conv2dStaticSamePadding(
32, 32, kernel_size=(3, 3), stride=[1, 1], groups=32, bias=False
(static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)
)
(_bn1): BatchNorm2d(32, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)
(_se_reduce): Conv2dStaticSamePadding(
32, 8, kernel_size=(1, 1), stride=(1, 1)
(static_padding): Identity()
)
(_se_expand): Conv2dStaticSamePadding(
8, 32, kernel_size=(1, 1), stride=(1, 1)
(static_padding): Identity()
)
...
现在有趣的是这个模型的最后几层:
...
(_bn1): BatchNorm2d(1280, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)
(_avg_pooling): AdaptiveAvgPool2d(output_size=1)
(_dropout): Dropout(p=0.2, inplace=False)
(_fc): Linear(in_features=1280, out_features=1000, bias=True)
(_swish): MemoryEfficientSwish()
假设我们想要重用这个模型主干，除了_fc，因为我们想在另一个具有相同主干但不同头部的模型上使用权重，而不是预先训练。在此示例中，我将采用相同的主干并添加 3 个头:
class ThreeHeadEfficientNet(torch.nn.Module):
def __init__(self,nbClasses1,nbClasses2,nbClasses3,model="efficientnet-b0",dropout_p=0.2):
super(ThreeHeadEfficientNet, self).__init__()
self.NBC1 = nbClasses1
self.NBC2 = nbClasses2
self.NBC3 = nbClasses3
self.dropout_p = dropout_p
self._dropout_layer = torch.nn.Dropout(p=self.dropout_p)
self._head1 = torch.nn.Linear(1280,self.NBC1)
self._head2 = torch.nn.Linear(1280,self.NBC2)
self._head3 = torch.nn.Linear(1280,self.NBC3)
self.model = EfficientNet.from_name(model,include_top=False) #you can notice here, I'm not loading the head, only the backbone
def forward(self,x):
features = self.model(x)
res = features.flatten(start_dim=1)
res = self._dropout_layer(res)
res1 = self._head1(res)
res2 = self._head2(res)
res3 = self._head3(res)
return res1,res2,res3
您现在会注意到，如果打印此 ThreeHeadsModel 图层，图层名称会从 _conv_stem.weight 稍微更改为 model._conv_stem.weight，因为主干是现在存储在属性变量model中。因此，我们必须处理这个问题，否则键将不匹配，创建一个新的状态字典来匹配这个新模型的预期键并包含预训练的权重和偏差:
pretrained_dict = model.state_dict() #pretrained model keys
model_dict = new_model.state_dict() #new model keys
processed_dict = {}
for k in model_dict.keys():
decomposed_key = k.split(".")
if("model" in decomposed_key):
pretrained_key = ".".join(decomposed_key[1:])
processed_dict[k] = pretrained_dict[pretrained_key] #Here we are creating the new state dict to make our new model able to load the pretrained parameters without the head.
new_model.load_state_dict(processed_dict, strict=False) #strict here is important since the heads layers are missing from the state, we don't want this line to raise an error but load the present keys anyway.
最后，在 new_model 中，您的新模型应该具有经过预训练的主干和需要微调的头部。
现在您应该能够解决您的问题了:)
更多pytorch信息，还请查看forum .
关于pytorch - 如何使用 PyTorch 保存和加载神经网络的特定层？，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/70986805/