代码之家 › 专栏 › 技术社区 › Farshid Rayhan

如何在稍加修改的模型中重用旧权重?

pytorch conv-neural-network

Farshid Rayhan · 技术社区 · 6 年前

我有一个这样的CNN网络是为一个特殊的任务而建立的。

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

该模型是使用这样的torch内置方法存储的。

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')

我在网络中增加了一层,而模型的其余部分保持不变。

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

由于其他conv层保持不变,是否可以重用保存的模型将权重加载到conv11、conv12和conv14?而不是从头开始训练?

1 回复 | 直到 6 年前

iacolippo 6 年前

假设你训练了下面的模型,现在你对它做了一个小的修改(比如添加一个层),并且想要使用你训练过的权重

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)

将模型(和优化器状态)保存为:

state = {'state_dict': net.state_dict(),
         'opt': optimizer.state_dict()
        }
torch.save(state, 'state.pt')

您的新模型是(请注意,相应的层保持相同的名称,因此您不会生成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

现在你可以加载 model.pt 文件:

state = torch.load('state.pt')

state 是一个 dict ,state['opt']包含优化器的所有参数,例如 state['opt']['param_groups'][0]['lr'] 给予

0.001

假设相应的层保持相同的名称,则可以通过以下方法恢复参数并初始化相应的层:

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])