Pytorch中实现只导入部分模型参数的方式

在PyTorch中，有时候我们只需要导入模型的部分参数，而不是全部参数。以下是两个示例说明，介绍如何在PyTorch中实现只导入部分模型参数的方式。

示例1：只导入部分参数

import torch
import torch.nn as nn

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.fc1 = nn.Linear(32 * 8 * 8, 10)
        self.fc2 = nn.Linear(10, 2)

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = x.view(-1, 32 * 8 * 8)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载模型
model = MyModel()
state_dict = torch.load('model.pth')

# 只导入部分参数
new_state_dict = {}
for k, v in state_dict.items():
    if 'conv' in k:
        new_state_dict[k] = v

# 更新模型参数
model.load_state_dict(new_state_dict)

在这个示例中，我们首先定义了一个名为MyModel的模型，并使用torch.load函数加载了一个名为model.pth的模型参数文件。然后，我们使用for循环遍历模型参数字典，并将包含conv的参数存储在一个新的字典中。最后，我们使用model.load_state_dict函数将新的参数字典加载到模型中。

示例2：只导入部分层的参数

import torch
import torch.nn as nn

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.fc1 = nn.Linear(32 * 8 * 8, 10)
        self.fc2 = nn.Linear(10, 2)

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = x.view(-1, 32 * 8 * 8)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载模型
model = MyModel()
state_dict = torch.load('model.pth')

# 只导入部分层的参数
new_state_dict = {}
for k, v in state_dict.items():
    if 'conv1' in k:
        new_state_dict[k] = v

# 更新模型参数
model.conv1.load_state_dict(new_state_dict)

在这个示例中，我们首先定义了一个名为MyModel的模型，并使用torch.load函数加载了一个名为model.pth的模型参数文件。然后，我们使用for循环遍历模型参数字典，并将包含conv1的参数存储在一个新的字典中。最后，我们使用model.conv1.load_state_dict函数将新的参数字典加载到模型的conv1层中。