I used these two functions to rename the features of a model:
def NameChanger(state_dict):
w, b, n = 0, 0, 0
new_state_dict = OrderedDict()
for k, v in state_dict.items():
if n < 10:
if "weight" in str(k):
basename = k[1:]
feature_name = "features." + str(w) + str(basename)
new_state_dict[feature_name] = v
print(feature_name)
w+=1
if "bias" in str(k):
basename = k[1:]
feature_name = "features." + str(b) + str(basename)
new_state_dict[feature_name] = v
print(feature_name)
b+=1
else:
if "weight" in str(k):
basename = k[2:]
feature_name = "features." + str(w) + str(basename)
new_state_dict[feature_name] = v
print(feature_name)
w+=1
if "bias" in str(k):
basename = k[2:]
feature_name = "features." + str(b) + str(basename)
new_state_dict[feature_name] = v
print(feature_name)
b+=1
n+=1
return new_state_dict
def ChangeNames(model_file):
state_dict = torch.load(model_file)
cnn = NIN_Conv(features)
new_state_dict = NameChanger(state_dict)
torch.save(new_state_dict, 'nin_imagenet_conv2.pth')
cnn.load_state_dict(new_state_dict)
return cnn
And this is the model definition:
import torch
import torch.nn as nn
from collections import OrderedDict
class NIN_Conv(nn.Module):
def __init__(self, features):
super(NIN_Conv, self).__init__()
self.features = features
def BuildSequential():
layers = []
conv2d = nn.Conv2d(3,96,(11, 11),(4, 4))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(96,96,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(96,96,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
pool2d = nn.MaxPool2d((3, 3),(2, 2),(0, 0),ceil_mode=True)
layers += [pool2d]
conv2d = nn.Conv2d(96,256,(5, 5),(1, 1),(2, 2))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(256,256,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(256,256,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
pool2d = nn.MaxPool2d((3, 3),(2, 2),(0, 0),ceil_mode=True)
layers += [pool2d]
conv2d = nn.Conv2d(256,384,(3, 3),(1, 1),(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(384,384,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(384,384,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
pool2d = nn.MaxPool2d((3, 3),(2, 2),(0, 0),ceil_mode=True)
layers += [pool2d]
drop2d = nn.Dropout(0.5)
layers += [drop2d]
conv2d = nn.Conv2d(384,1024,(3, 3),(1, 1),(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(1024,1024,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
conv2d = nn.Conv2d(1024,1000,(1, 1))
layers += [conv2d, nn.ReLU(inplace=True)]
pool2d = nn.AvgPool2d((6, 6),(1, 1),(0, 0),ceil_mode=True)
layers += [pool2d]
soft2d = nn.Softmax()
layers += [soft2d]
return nn.Sequential(*layers)
When I try to load my model, with this:
cnn = NIN_Conv(BuildSequential())
cnn.load_state_dict(torch.load('nin_imagenet_conv2.pth'))
I get the following error:
cnn.load_state_dict(torch.load('nin_imagenet_conv2.pth'))
File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 522, in load_state_dict
.format(name))
KeyError: 'unexpected key "features.1.weight" in state_dict'
The same errror also occurs inside the ChangeNames function, at:
cnn.load_state_dict(new_state_dict)
How do I resolve this error?