Load convolutional layers partially from saved model

vision
1

Is there a way to load specific layers of convolutional layers (e.g. only 1-3), when I saved the model in following way:

torch.save({'state_dict': net.state_dict(),'optimizer': optimizer.state_dict()},PATH)

instead of

torch.save(net.state_dict(), PATH)

My network looks like:

        self.conv.add_module('conv1_s1',nn.Conv2d(3, 32, kernel_size=9, stride=1, padding=0))
        self.conv.add_module('relu1_s1',nn.ReLU(inplace=True))
        self.conv.add_module('pool1_s1',nn.MaxPool2d(kernel_size=3, stride=1))                           

        self.conv.add_module('conv2_s1',nn.Conv2d(32, 48, kernel_size=3, padding=2, groups=2))
        self.conv.add_module('relu2_s1',nn.ReLU(inplace=True))
        self.conv.add_module('pool2_s1',nn.MaxPool2d(kernel_size=3, stride=2))

        self.conv.add_module('conv3_s1',nn.Conv2d(48, 96, kernel_size=3, padding=2, groups=2))
        ..
        self.conv.add_module('conv4_s1',nn.Conv2d(96, 96, kernel_size=3, padding=2, groups=2))
        ..
        self.conv.add_module('conv5_s1',nn.Conv2d(96, 48, kernel_size=3, padding=2, groups=2))

I tried to load the layer partially from saved model like:

        checkpointL = torch.load(PATH)
        originalNetwork.load_state_dict(checkpointL['conv.conv1_s1.weight'])
        originalNetwork.load_state_dict(checkpointL['conv.conv1_s1.bias'])

and got error: KeyError: 'conv.conv1_s1.weight'

Do I have to save the model using torch.save(net.state_dict(), PATH) in order to be able to load only specific layers of my convolutional layers.

Thanks for help!

2

I think the way you have saved the model seems fine
Can you print checkpointL.keys() values?

3

checkpointL.keys() yields dict_keys(['state_dict', 'optimizer']). I want to load the weights for only e.g. 'conv1_s1' or 'conv2_s1'.