pytorch 提取权重_PyTorch中如何加载子模块的权重

假设我们在深度学习模型中有一个这样的需求：主要模型A中包含子模块B，而模型B可以通过一定的方式得到一个预训练的权重，模型A需要利用B模型的权重，在此基础上继续训练。

首先我们到官网上去寻找，PyTorch官网上给出了一些保存和加载模型的示例，可以说非常全面总结了模型保存和加载的方法和主义事项，https://pytorch.org/tutorials/beginner/saving_loading_models.html。但是这里的方案都是针对一个完整模型的保存和加载的，不能满足我们这个需求。

因此需要基于此做一些改进，具体如代码所示：

import torch

import torch.nn as nn

class ModelA(nn.Module):

def __init__(self):

super(ModelA, self).__init__()

self.A = nn.Linear(2, 3)

def forward(self, A):

pass

class ModelB(nn.Module):

def __init__(self):

super(ModelB, self).__init__()

self.model_a = ModelA()

self.A = nn.Linear(2, 3)

def forward(self, x):

pass

print("Model")

modelA = ModelA()

modelA_dict = modelA.state_dict()

print('-' * 80)

for key in sorted(modelA_dict.keys()):

parameter = modelA_dict[key]

print(key)

print(parameter.size())

print(parameter)

modelB = ModelB()

modelB_dict = modelB.state_dict()

print('-'*80)

for key in sorted(modelB_dict.keys()):

print('-'*20)

parameter = modelB_dict[key]

print(type(key), key)

print(parameter.size())

print(parameter)

print('-'*20)

print('-'*80)

pretrained_dict = modelA_dict

model_dict = modelB_dict

pretrained_dict = {'model_a.' + k: v for k, v in pretrained_dict.items() if 'model_a.' + k in model_dict}

model_dict.update(pretrained_dict)

modelB.load_state_dict(model_dict)

modelB_dict = modelB.state_dict()

for key in sorted(modelB_dict.keys()):

parameter = modelB_dict[key]

print(key)

print(parameter.size())

print(parameter)