如何使用PyTorch保存和加载特定层次的神经网络？
裸问题语句
我已经训练了一个模型A，其中包括一个特征提取FE和分类头ACH。
我想训练一个B型，它使用A的特征提取器FE和再训练它自己的分类头BCH。
到目前为止还很简单。现在我不想保存整个模型B，因为它的FE部分已经保存在模型A中。我只想转储BCH，在推理期间
- 负荷模型A-做它的预测
- 负荷B的分类头BCH。
- 用BCH交换分类头ACH
- 使用此交换状态运行预测。
阅读pyTorches文档，它只提到保存整个模型。我怎样才能做到这一点？
问题陈述的结束
更多关于问题动机的详细信息：
我有一个图像数据集，我想分类，这些图像可以有几个类给他们。例如，同一图像可以有“陆地车辆”(超级类别)和一类“汽车”(类别)或“卡车”类。另一幅图像可能是“飞行器”级，它可以是“直升机”或“飞机”。
由于图像和大部分特征应该是相同的，我希望为超类别训练一个分类器，然后冻结它的特征提取器，然后使用预先训练的特征提取器来学习相同的类别模型。
由于特征提取的主干网的权重是相同的，所以我只想保存分类头的权重，从而节省一些宝贵的计算资源。
回答 1
Stack Overflow用户
发布于 2022-02-04 14:44:04
一般来说，通常只需要访问模型的主干，以便将其用于其他目的。您有几种方法来执行此操作。但是，大多数情况下，考虑到保存模型检查点并在以后加载它意味着保存权重和偏差，并且能够正确地将它们加载到相应的层，您首先需要从模型中了解您想要保存的部分。
当您获得模型的状态时，您将获得一个字典。键是层名，值是权重和偏差。让我们看一个关于如何只保存模型的主干的效率网分类器示例。基本上，效率网，如你的例子，是一个主干网和一个完全连接的层作为一个头，如果你只想要骨干，你想要每一层，除了头部，你将微调以后。
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层，由于主干现在存储在属性变量model中，所以层名从_conv_stem.weight略微更改为model._conv_stem.weight。因此，我们必须处理否则键将不匹配的问题，创建一个新的状态字典，该字典与这个新模型的预期键匹配，并包含预先训练过的权重和偏差：
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中，您应该有您的新模型与预先培训的脊梁和头部的微调。
现在您应该能够解决您的问题了:)
欲知更多火把的信息，请也检查论坛。
https://stackoverflow.com/questions/70986805
复制相似问题
腾讯云开发者
