Question

我为特定任务建立了这样的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

使用类似的割炬内置方法存储模型。

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？而不是从头开始训练？

Answer 1

假设您训练了以下模型，现在您对其进行了较小的修改（例如添加图层），并希望使用训练后的权重

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

如何在经过稍微修改的模型中重用旧的权重？

1 个答案: