python - 如何在 Pytorch 中的模块之间共享权重？

Question

在 Pytorch 中的两层（模块）之间共享权重的正确方法是什么？
根据我在 Pytorch 论坛上的发现，有几种方法可以做到这一点。
例如，基于这个讨论，我认为简单地分配转置权重就可以了。那是在做：

 self.decoder[0].weight = self.encoder[0].weight.t()

然而，这被证明是错误的并导致错误。然后我尝试将上面的行包装在nn.Parameter()：

self.decoder[0].weight = nn.Parameter(self.encoder[0].weight.t())

这消除了错误，但话又说回来，这里没有发生共享。通过这个，我刚刚初始化了一个新的张量，其值与encoder[0].weight.t().

然后我找到了这个链接，它提供了不同的权重共享方式。但是，我怀疑那里给出的所有方法实际上是否正确。
例如，一种方法如下所示：

# tied autoencoder using off the shelf nn modules
class TiedAutoEncoderOffTheShelf(nn.Module):
    def __init__(self, inp, out, weight):
        super().__init__()
        self.encoder = nn.Linear(inp, out, bias=False)
        self.decoder = nn.Linear(out, inp, bias=False)

        # tie the weights
        self.encoder.weight.data = weight.clone()
        self.decoder.weight.data = self.encoder.weight.data.transpose(0,1)

    def forward(self, input):
        encoded_feats = self.encoder(input)
        reconstructed_output = self.decoder(encoded_feats)
        return encoded_feats, reconstructed_output

基本上，它使用创建一个新的权重张量nn.Parameter()并将其分配给每个层/模块，如下所示：

weights = nn.Parameter(torch.randn_like(self.encoder[0].weight))
self.encoder[0].weight.data = weights.clone()
self.decoder[0].weight.data = self.encoder[0].weight.data.transpose(0, 1)

这真的让我很困惑，这如何在这两层之间共享相同的变量？不只是克隆“原始”数据吗？
当我使用这种方法并可视化权重时，我注意到可视化是不同的，这让我更加确定有些事情是不对的。
我不确定不同的可视化是否仅仅是由于一个是另一个的转置，或者正如我已经怀疑的那样，它们是独立优化的（即权重不在层之间共享）

示例权重初始化：

score 2 · Accepted Answer

艾问题一般容易被错误理解，尤其是这个问题。我将改写您的问题，模块 M1 的 A 层和模块 M2 的 B 层可以共享权重 WA = WB，并且可能的 WA = WB_transposed。

这可以通过 PyTorch 钩子实现，您可以在其中更新 A 的前向钩子以更改 WB，并且您可能会在 M2 autograd 中冻结 WB。

所以只需使用钩子。

from time import sleep

import torch
import torch.nn as nn
class M(nn.Module):
    def __init__(self):        
        super().__init__()        
        self.l1 = nn.Linear(1,2)

    def forward(self, x):                      
        x = self.l1(x)
        return x

model = M()
model.train()

def printh(module, inp, outp):
    sleep(1)    
    print("update other model parameter in here...")   


h = model.register_forward_hook(printh)
for i in range(1,4):

    x = torch.randn(1)
    output = model(x)

h.remove()

score 1 · Accepted Answer

有趣的是，您对@Rika 的第一个直觉是正确的：

这真的让我很困惑，这如何在这两层之间共享相同的变量？不只是克隆“原始”数据吗？

很多人实际上在博客文章或他们自己的存储库中都犯了这个错误。

正如您所写 self.decoder[0].weight = nn.Parameter(self.encoder[0].weight.t()) ，还将简单地创建一个新的权重矩阵。

唯一可行的做法似乎是使用由 nn.Linear ( torch.nn.functional.linear()) 调用的线性函数：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init

# Real off-the-shelf tied linear module
class TiedLinear(nn.Module):
    def __init__(self, tied_to: nn.Linear, bias: bool = True):
        super().__init__()
        self.tied_to = tied_to
        if bias:
            self.bias = nn.Parameter(torch.Tensor(tied_to.in_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    # copied from nn.Linear
    def reset_parameters(self):
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.tied_to.weight.t())
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return F.linear(input, self.tied_to.weight.t(), self.bias)

    # To keep module properties intuitive
    @property
    def weight(self) -> torch.Tensor:
        return self.tied_to.weight.t()

# Shared weights, different biases
encoder = nn.Linear(in, out)
decoder = TiedLinear(encoder)

score 1 · Accepted Answer

事实证明，经过进一步调查，这只是简单地重新转换解码器的权重并将其可视化，它们确实是共享的。
下面是编码器和解码器权重的可视化：

python - 如何在 Pytorch 中的模块之间共享权重？

3 回答 3

Related

Reference