How to re-use old weights in a slightly modified model? - conv-neural-network

Question

I have a CNN network built like this for a particular task.
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv11 = nn.Conv2d(1, 128, kernel_size=3, padding=1)
self.conv12 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
self.conv13 = nn.Conv2d(256, 2, kernel_size=3, padding=1)
def forward(self, x):
in_size = x.size(0)
x = F.relu(self.conv11(x))
x = F.relu(self.conv12(x))
x = F.relu(self.conv13(x))
x = F.softmax(x, 2)
return x
The model is stored using the torch built-in method like this.
net = Net()
optimizer = optim.SGD(net.parameters(), lr=1e-3)
state = {
'state_dict': net.state_dict()
'opt': optimizer.state_dict()
}
torch.save(state, 'model.pt')
I have increased a single layer in the network while the rest of the model was kept the same.
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv11 = nn.Conv2d(1, 128, kernel_size=3, padding=1)
self.conv12 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
self.conv13 = nn.Conv2d(256, 256, kernel_size=3, padding=1) # (new added)
self.conv14 = nn.Conv2d(256, 2, kernel_size=3, padding=1)
def forward(self, x):
in_size = x.size(0)
x = F.relu(self.conv11(x))
x = F.relu(self.conv12(x))
x = F.relu(self.conv13(x)) (new added)
x = F.relu(self.conv14(x))
x = F.softmax(x, 2)
return x
Since the other conv layers are kept the same, is there any way I can re-use the saved model to load the weights to conv11, conv12 and conv14 ? Instead of starting to train from beginning ?

Answer 1

Assume you trained the following model and now you make a minor modification to it (like adding a layer) and want to use your trained weights
import torch
import torch.nn as nn
import torch.optim as optim
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv11 = nn.Conv2d(1, 128, kernel_size=3, padding=1)
self.conv12 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
self.conv13 = nn.Conv2d(256, 2, kernel_size=3, padding=1)
def forward(self, x):
in_size = x.size(0)
x = F.relu(self.conv11(x))
x = F.relu(self.conv12(x))
x = F.relu(self.conv13(x))
x = F.softmax(x, 2)
return x
net = Net()
optimizer = optim.SGD(net.parameters(), lr=1e-3)
you save the model (and the optimizer state) with:
state = {'state_dict': net.state_dict(),
'opt': optimizer.state_dict()
}
torch.save(state, 'state.pt')
Your new model is (note that corresponding layers keep the same name, so you don't make conv13 -> conv14):
class NewNet(nn.Module):
def __init__(self):
super(NewNet, self).__init__()
self.conv11 = nn.Conv2d(1, 128, kernel_size=3, padding=1)
self.conv12 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
self.convnew = nn.Conv2d(256, 256, kernel_size=3, padding=1) # (new added)
self.conv13 = nn.Conv2d(256, 2, kernel_size=3, padding=1)
def forward(self, x):
in_size = x.size(0)
x = F.relu(self.conv11(x))
x = F.relu(self.conv12(x))
x = F.relu(self.convnew(x)) # (new added)
x = F.relu(self.conv13(x))
x = F.softmax(x, 2)
return x
Now you can load your model.pt file:
state = torch.load('state.pt')
state is a dict, state['opt'] contains all the parameters that you had for your optimizer, for example state['opt']['param_groups'][0]['lr'] gives
0.001
Assuming corresponding layers kept the same name, you can recover your parameters and initialize the appropriate layers by:
net = NewNet()
for name, param in net.named_parameters():
if name in state['state_dict'].keys():
param = param.data
param.copy_(state['state_dict'][name])