KeyError when loading state_dict into my model

1

Hi, I have parameters param.pth and I want to load them into model like:

model = Net()
model.load_state_dict(torch.load("param.pth"))

but it cause

KeyError: 'unexpected key "0.weight" in state_dict'

How can I load params in state_dict?

Following are details of Net and param.pth:

Network

class Net(nn.Module):
    def __init__(self):
        super(StyleNet, self).__init__()
        self.conv1 = nn.Conv2d(3,64,(3, 3),(1, 1),(1, 1))
        self.conv2 = nn.Conv2d(64,64,(3, 3),(1, 1),(1, 1))
        self.conv2_drop = nn.Dropout(0.25)
        self.pool1 = nn.MaxPool2d((4, 4),(4, 4))
        self.bn1 = nn.BatchNorm2d(64,0.001,0.9,True)
        self.conv3 = nn.Conv2d(64,128,(3, 3),(1, 1),(1, 1))
        self.conv4 = nn.Conv2d(128,128,(3, 3),(1, 1),(1, 1))
        self.conv4_drop = nn.Dropout(0.25),
        self.pool2 = nn.MaxPool2d((4, 4),(4, 4))
        self.bn2 = nn.BatchNorm2d(128,0.001,0.9,True)
        self.conv5 = nn.Conv2d(128,256,(3, 3),(1, 1),(1, 1))
        self.conv6 = nn.Conv2d(256,256,(3, 3),(1, 1),(1, 1))
        self.conv6_drop = nn.Dropout(0.25)
        self.pool3 =nn.MaxPool2d((4, 4),(4, 4))
        self.bn3 = nn.BatchNorm2d(256,0.001,0.9,True)
        self.conv7 = nn.Conv2d(256,128,(3, 3),(1, 1),(1, 1))
        self.linear1 = nn.Linear(3072,128)

    def forward(self, input):
        x = nn.ReLU(self.conv1(input))
        x = nn.ReLU(self.conv2(input))
        x = self.conv2_drop(x)
        x = self.bn1(self.pool1(x))
        x = nn.ReLU(self.conv3(x))
        x = self.conv4_drop(nn.ReLU(self.conv4(x)))
        x = self.bn2(self.pool2(x))
        x = nn.ReLU(self.conv5(x))
        x = nn.ReLU(self.conv6(x))
        x = self.conv6_drop(x)
        x = self.bn3(self.pool3(x))
        x = nn.ReLU(self.conv7(x))

        return self.linear1(x)

Keys of params.pth

in

torch.load("params.pth").keys()

out

odict_keys(['0.weight', '0.bias', '2.weight', '2.bias', '6.weight', '6.bias', '6.running_mean', '6.running_var', '7.weight', '7.bias', '9.weight', '9.bias', '13.weight', '13.bias', '13.running_mean', '13.running_var', '14.weight', '14.bias', '16.weight', '16.bias', '20.weight', '20.bias', '20.running_mean', '20.running_var', '21.weight', '21.bias', '24.1.weight', '24.1.bias'])

Thank you for dealing with it.

2

It looks like you have used nn.Sequential for saving the state_dict.
Could you try to wrap all your layers into nn.Sequential and load it again?

3

Actually, the model was saved on nn.Sequential form, so I can load them when I defined the model as nn.Sequential.
Is there any simple way to save the parameters so that model defined above can load them?
Thanks.

4

Well, there is a way, but it’s a bit ugly, so sorry for that, but I can’t come up with a cleaner way at the moment.
I’ve created a small example how to change the keys in the state_dict:

# Create Sequential model
modelA = nn.Sequential(
    nn.Conv2d(3, 6, 3, 1, 1),
    nn.ReLU(),
    nn.Conv2d(6, 12, 3, 1, 1),
    nn.LogSoftmax(dim=1))

# Create Module
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6 ,3, 1, 1)
        self.conv2 = nn.Conv2d(6, 12 ,3, 1, 1)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.conv2(x)
        return F.log_softmax(x, dim=1)


modelB = MyModel()

from collections import OrderedDict

# Get Sequential state dict
state_dict = modelA.state_dict()
for keyA, keyB in zip(modelA.state_dict(), modelB.state_dict()):
    print('Changing {} to {}'.format(keyA, keyB))
    state_dict = OrderedDict((keyB if k == keyA else k, v) for k, v in state_dict.items())

# state_dict should keep the old values with new keys
modelB.load_state_dict(state_dict)
print(modelB.state_dict())

Could you try that?

2 Likes
5

Hi, @ptrblck

I modified your code to my case, and it properly worked!
I really thank you!