加载预训练模型的部分参数

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

前言

例如：如果你想用已经训练好的模型来作为你的任务的特征提取器，可是不需要完整的模型，模型的一部分提取特征就可以，那么看我下面的例子，分两步来详细说明如何加载部分模型的参数。

---

第一步：按参数文件实现模型

下面我用C3D模型来做例子，C3D模型的实现代码，我只用了fc6层之前的模型部分，将后面的层都注释掉了。

import torch
import torch.nn as nn

class C3D(nn.Module):

    def __init__(self):
        super(C3D, self).__init__()

        self.conv1 = nn.Conv3d(3, 64, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool1 = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))

        self.conv2 = nn.Conv3d(64, 128, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool2 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))

        self.conv3a = nn.Conv3d(128, 256, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.conv3b = nn.Conv3d(256, 256, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool3 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))

        self.conv4a = nn.Conv3d(256, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.conv4b = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool4 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))

        self.conv5a = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.conv5b = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool5 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0, 1, 1))

        self.fc6 = nn.Linear(8192, 4096)
        # self.fc7 = nn.Linear(4096, 4096)
        # self.fc8 = nn.Linear(4096, 487)

        # self.dropout = nn.Dropout(p=0.5)

        self.relu = nn.ReLU()
        # self.softmax = nn.Softmax(dim=1)

    def forward(self, x):

        h = self.relu(self.conv1(x))
        h = self.pool1(h)

        h = self.relu(self.conv2(h))
        h = self.pool2(h)

        h = self.relu(self.conv3a(h))
        h = self.relu(self.conv3b(h))
        h = self.pool3(h)

        h = self.relu(self.conv4a(h))
        h = self.relu(self.conv4b(h))
        h = self.pool4(h)

        h = self.relu(self.conv5a(h))
        h = self.relu(self.conv5b(h))
        h = self.pool5(h)
        h = h.view(-1, 8192)
        h = self.relu(self.fc6(h))
        # h = self.dropout(h)
        # h = self.relu(self.fc7(h))
        # h = self.dropout(h)

        # logits = self.fc8(h)
        # probs = self.softmax(logits)

        return h

第二步：加载参数

C3D 模型的参数文件‘c3d.pickle’，我已经提前下载并保存在本地了。

# loading the altered C3D (ie C3D upto before fc-6)
model_pretrained_dict = torch.load('c3d.pickle')  #加载预训练模型的全部参数
model = C3D() #实例化模型
model_dict = C3D.state_dict()  #将实例化模型的全部参数读出
model_pretrained_dict = {k:v for k, v in model_pretrained_dict.items() if k in mode_dict}  #只留下模型需要的那部分的参数
model_dict.update(model_pretrained_dict)  #将模型开始的随机参数替换为预训练参数
model.load_state_dict(model_dict)  #将参数加载进模型
model = model.cuda()    #将模型加载到GPU中