RuntimeError when loading model

vision
1

Hello,

I got the following error when I tried to load my model:

RuntimeError: Error(s) in loading state_dict for VGG:

Missing key(s) in state_dict: “classifier.0.weight”, “classifier.0.bias”, “classifier.3.weight”, “classifier.3.bias”, “classifier.6.weight”, “classifier.6.bias”.

Unexpected key(s) in state_dict: “classifier.fc1.weight”, “classifier.fc1.bias”, “classifier.fc2.weight”, “classifier.fc2.bias”.

The code that I used to save and load my model is as follow:

torch.save(model.state_dict(), ‘check_point.pth’)
model.load_state_dict(torch.load(‘check_point.pth’))

What am I missing?

2

Are you using now the nn.Sequential wrapper for your classifier?
It looks like you’ve saved your model using layers fc1 and fc2 while these layers are now wrapped in nn.Sequential.
If so, you could try to use an OrderedDict to set the layer names again:

self.classifier = nn.Sequential(OrderedDict([
          ('fc1', nn.Linear(...)),
          ('relu1', nn.ReLU()),
          ('fc2', nn.Linear(...)),
          ('relu2', nn.ReLU())
]))
3

Hello @ptrblck,

Thanks for your quick reply!

I have used the solution as you suggested but I am still facing the same error. My codes as follow:

Build Classifier

model = models.vgg19(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

classifier = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(25088, 4096)),
    ('relu1', nn.ReLU()),
    ('dropout1', nn.Dropout()),
    ('fc2', nn.Linear(4096, 102)),
    ('output', nn.LogSoftmax(dim=1))
]))

model.classifier = classifier    
model.to(device)

criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=0.001)

Saving and Loading Model

torch.save(model.state_dict(), 'check_point.pth')

model.load_state_dict(torch.load('check_point.pth'))

However, I still receive the same RunTime Error when I tried to load the check_point.

On a side note, there is a tutorial from PyTorch on saving and loading model (https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). The example seems to be for building a new CNN rather than using an existing pre-trained model?

Below is the code that is illustrated:

# Define model
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Initialize model
model = TheModelClass()

My Questions:

  1. What am I missing in my code that is causing the RunTime Error?
  2. If I am using a pre-trained model, do I still need to define TheModelClass?
  3. Is the following line of code necessary if I am using a pre-trained model?
    the_model = TheModelClass(*args, **kwargs)

Thanks!

4

In your code example you are using a pretrained VGG model and change the classifier.
After the training you are saving the state_dict and try to load it afterwards.
While loading the state_dict you are experiencing the error. Is this correct?

Could you explain which model you’ve saved and how you’ve saved it?
I don’t get any errors using your code and the error message points to a mismatch between the keys in the state_dict. It seems your model uses keys from nn.Sequential (e.g. classifier.0.weight) while the stored state_dict has named keys e.g. classifier.fc1.weight.

TheModelClass is just a dummy model to show, how to store and load the state_dict.
You need an instance of your model to load the checkpoint.
This code works fine:

# Store model
model = models.vgg19(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

classifier = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(25088, 4096)),
    ('relu1', nn.ReLU()),
    ('dropout1', nn.Dropout()),
    ('fc2', nn.Linear(4096, 102)),
    ('output', nn.LogSoftmax(dim=1))
]))

model.classifier = classifier
torch.save(model.state_dict(), 'vgg_test.pth')

# Load model
model = models.vgg19(pretrained=False)
for param in model.parameters():
    param.requires_grad = False

classifier = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(25088, 4096)),
    ('relu1', nn.ReLU()),
    ('dropout1', nn.Dropout()),
    ('fc2', nn.Linear(4096, 102)),
    ('output', nn.LogSoftmax(dim=1))
]))
model.classifier = classifier
model.load_state_dict(torch.load('vgg_test.pth'))

If you comment out the model.classifier = classifier line before loading the state_dict, you’ll get the same error you are describing.

2 Likes