pytorch 复制、粘贴、替换Model某一层的训练参数

使用pytorch保存模型参数与加载模型参数的方法可以对某一层的参数进行替换、复制：

打印当前网络的参数：nn.Model.state_dict()

加载网络参数：nn.Model.load_state_dict()

代码示例

mport torch.nn as nn


class Net1(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear1 = nn.Linear(input_dim, output_dim)  # 输入的个数，输出的个数
        self.linear2 = nn.Linear(output_dim, 2)

    def forward(self, x):
        out = self.linear1(x)
        out = self.linear2(out)
        return out


class Net2(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)  # 输入的个数，输出的个数

    def forward(self, x):
        out = self.linear(x)
        return out


if __name__ == '__main__':
    net1 = Net1(20, 2) # 模型1
    net2 = Net2(20, 2) # 模型2
    # 这里可以打断点，debug一下确认参数是否变更哟
    net2.linear.load_state_dict(net1.linear1.state_dict()) # 把模型2的fc层参数替换为模型1的linear1参数

更多《计算机视觉与图形学》知识，可关注下方公众号：

参考文章

pytorch获取模型某一层参数名及参数值方式：https://www.cnblogs.com/expttt/p/12461775.html  
In Pytorch, what is the most efficient way to copy the learned params of a model as the initialization for a second model of the same architecture?：https://stackoverflow.com/questions/59183865/in-pytorch-what-is-the-most-efficient-way-to-copy-the-learned-params-of-a-model