从预训练模型加载单层权重

Question

我想专门将一些层的预训练模型参数添加到我的新网络中。 对于我刚刚做的线性层：

model_enc.linear_3d.weight = model_trained.linear_3d.weight
model_enc.linear_3d.bias = model_trained.linear_3d.bias

这是否足够，或者是否有任何其他我需要加载的参数，或者有比这更简单的方法。我的模块已经训练过了，我只想为几层加载参数。

谢谢

Answer 1

您的解决方案应该可行，而且对我来说似乎很简单。

从https://pytorch.org/docs/master/_modules/torch/nn/modules/linear.html#Linear上的源码可以看出，nn.Linear模块有in_features、out_features、weight1和bias属性：

def __init__(self, in_features, out_features, bias=True):
    super(Linear, self).__init__()
    self.in_features = in_features
    self.out_features = out_features
    self.weight = Parameter(torch.Tensor(out_features, in_features))
    if bias:
        self.bias = Parameter(torch.Tensor(out_features))
    else:
        self.register_parameter('bias', None)
    self.reset_parameters()

因此，只要您的 in_features 和 out_features 相同，您就可以像以前一样替换权重和偏差。

或者，如果您将其存储为属性，您可以将一个网络中的整个线性模块替换为另一个网络的模块。