如何解决pytorch不会保存加载的预训练模型权重及其部分内容到最终模型中
我目前正在根据数据在CIFAR-10上进行预训练的模型，删除了模型的最后fc层，并附加了我自己的fc层和softmax。有七个网络，每个网络都与预训练部分相同，并使用附加的fc层进行组合。以下是经过预先训练的网络代码：
class Bottleneck(nn.Module):
def __init__(self,inplanes,expansion=4,growthRate=12,dropRate=0):
super(Bottleneck,self).__init__()
planes = expansion * growthRate
self.bn1 = nn.BatchNorm2d(inplanes)
self.conv1 = nn.Conv2d(inplanes,planes,kernel_size=1,bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes,growthRate,kernel_size=3,padding=1,bias=False)
self.relu = nn.ReLU(inplace=True)
self.dropRate = dropRate
def forward(self,x):
out = self.bn1(x)
out = self.relu(out)
out = self.conv1(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv2(out)
if self.dropRate > 0:
out = F.dropout(out,p=self.dropRate,training=self.training)
out = torch.cat((x,out),1)
return out
class BasicBlock(nn.Module):
def __init__(self,expansion=1,dropRate=0):
super(BasicBlock,x):
out = self.bn1(x)
out = self.relu(out)
out = self.conv1(out)
if self.dropRate > 0:
out = F.dropout(out,1)
return out
class Transition(nn.Module):
def __init__(self,outplanes):
super(Transition,self).__init__()
self.bn1 = nn.BatchNorm2d(inplanes)
self.conv1 = nn.Conv2d(inplanes,outplanes,bias=False)
self.relu = nn.ReLU(inplace=True)
def forward(self,x):
out = self.bn1(x)
out = self.relu(out)
out = self.conv1(out)
out = F.avg_pool2d(out,2)
return out
class DenseNet(nn.Module):
def __init__(self,depth = 22,block = Bottleneck,dropRate = 0,num_classes = 10,growthRate = 12,compressionRate = 2):
super(DenseNet,self).__init__()
assert (depth - 4) % 3 == 0,'depth should be 3n+4'
n = (depth - 4) / 3 if block == BasicBlock else (depth - 4) // 6
self.growthRate = growthRate
self.dropRate = dropRate
# self.inplanes is a global variable used across multiple
# helper functions
self.inplanes = growthRate * 2
self.conv1 = nn.Conv2d(3,self.inplanes,kernel_size = 3,padding = 1,bias = False)
self.dense1 = self._make_denseblock(block,n)
self.trans1 = self._make_transition(compressionRate)
self.dense2 = self._make_denseblock(block,n)
self.trans2 = self._make_transition(compressionRate)
self.dense3 = self._make_denseblock(block,n)
self.bn = nn.BatchNorm2d(self.inplanes)
self.relu = nn.ReLU(inplace=True)
self.avgpool = nn.AvgPool2d(8)
#self.fc = nn.Linear(self.inplanes,num_classes)
# Weight initialization
# for m in self.modules():
# if isinstance(m,nn.Conv2d):
# n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
# m.weight.data.normal_(0,math.sqrt(2. / n))
# elif isinstance(m,nn.BatchNorm2d):
# m.weight.data.fill_(1)
# m.bias.data.zero_()
def _make_denseblock(self,block,blocks):
layers = []
for i in range(blocks):
# Currently we fix the expansion ratio as the default value
layers.append(block(self.inplanes,growthRate = self.growthRate,dropRate=self.dropRate))
self.inplanes += self.growthRate
return nn.Sequential(*layers)
def _make_transition(self,compressionRate):
inplanes = self.inplanes
outplanes = int(math.floor(self.inplanes // compressionRate))
self.inplanes = outplanes
return Transition(inplanes,outplanes)
def forward(self,x):
x = self.conv1(x)
x = self.trans1(self.dense1(x))
x = self.trans2(self.dense2(x))
x = self.dense3(x)
x = self.bn(x)
x = self.relu(x)
x = self.avgpool(x)
#x = x.view(x.size(0),-1)
#x = self.fc(x)
return x
def getParams(self,paramName):
if paramName == 'inplanes':
return self.inplanes
elif paramName == 'growthRate':
return self.growthRate
elif paramName == 'dropRate':
return self.dropRate
def densenet(**kwargs):
"""
Constructs a DenseNet model.
"""
return DenseNet(**kwargs)
下面是我的代码：
class Network(nn.Module):
def __init__(self,pretrained_dict,num_classes = 6,num_channels = 7,expansion = 4,depth = 100,dropRate = 0):
super(Network,self).__init__()
self.num_channels = num_channels
# creating 7 channels networks
self.channels_dnsnets = []
for ch in range(self.num_channels):
# print(ch)
d = densenet(depth = depth)
d_dict = d.state_dict()
# 1. filter out unnecessary keys
pretrained_dict2 = {k[7:]: v for k,v in pretrained_dict.items() if k[7:] in d_dict}
# print('d_dict_keys :')
# print(d_dict.keys())
# print('*'*50)
# print('pretrained_dict2.keys:')
# print(pretrained_dict2.keys())
# print('*'*50)
# 2. overwrite entries in the existing state dict
d_dict.update(pretrained_dict2)
# 3. load the new state dict
d.load_state_dict(pretrained_dict2)
# freeze the layers of densenet
for param in d.parameters():
param.requires_grad = False
self.channels_dnsnets.append(d)
self.inplanes = self.channels_dnsnets[0].getParams(paramName = 'inplanes')
self.fc = nn.Linear(self.inplanes * self.num_channels,num_classes)
self.softmax = nn.Softmax(dim = 1)
def forward(self,x):
batch_size,channels,ht,wd,in_channels = x.shape
x = np.reshape(x,(batch_size,in_channels,wd))
out = []
for num in range(self.num_channels):
temp_out = self.channels_dnsnets[0](x[:,num,:])
temp_out = temp_out.view(temp_out.size(0),-1)
# print(temp_out.shape)
# print('*' * 50)
out.append(temp_out)
out = torch.stack(out,dim = 1)
# print(out.shape)
out = out.view(out.size(0),-1)
out = self.fc(out)
out = self.softmax(out)
return out
我将优化器设置为：
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,model.parameters()),lr = lr,betas = (0.9,0.999),eps = 1e-08,weight_decay = wd,amsgrad = False)
但是，每当我保存模型时，密集网列表及其权重都不会保存，而只会保存fc层和softmax层权重。代码有什么问题吗？我是pytorch的新手。
解决方法
问题是self.channels_dnsnets只是list，不会成为state_dict的一部分。仅self.fc和self.softmax将被注册到Module中。最简单的更改就是这样定义：
self.channels_dnsnets = nn.ModuleList()
版权声明：本文内容由互联网用户自发贡献，该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容， 请发送邮件至 dio@foxmail.com 举报，一经查实，本站将立刻删除。