python - 在 Pytorch 中创建自定义连接/非全连接层

Question

如图所示，它是一个带有NN的3层，分别是输入层、隐藏层和输出层。我想设计 NN(在 PyTorch 中，只是拱门)，其中隐藏层的输入是完全连接的。然而，从隐藏层到输出，隐藏层的前两个神经元应该连接到输出层的第一个神经元，后两个应该连接到输出层的第二个神经元，依此类推。这应该如何设计？

from torch import nn
layer1 = nn.Linear(input_size, hidden_size)
layer2 = ??????

Answer 1

作为@Jan说here , 你可以重载 nn.Linear并提供一个逐点掩码来掩蔽您想要避免的交互。请记住，全连接层只是带有可选加性偏差的矩阵乘法。

查看其 source code ，我们可以这样做:

class MaskedLinear(nn.Linear):
    def __init__(self, *args, mask, **kwargs):
        super().__init__(*args, **kwargs)
        self.mask = mask

    def forward(self, input):
        return F.linear(input, self.weight, self.bias)*self.mask

将 F 定义为 torch.nn.functional

考虑到您对第二层的约束:

the first two neurons of the hidden layer should be connected to the first neuron of the output layer

看来您正在寻找这种模式:

tensor([[1., 0., 0.],
        [1., 0., 0.],
        [0., 1., 0.],
        [0., 1., 0.],
        [0., 0., 1.],
        [0., 0., 1.]])

可以使用 torch.block_diag 获得:

mask = torch.block_diag(*[torch.ones(2,1),]*output_size)

有了这个，您可以将您的网络定义为:

net = nn.Sequential(nn.Linear(input_size, hidden_size),
                    MaskedLinear(hidden_size, output_size, mask))

---

如果你喜欢，你甚至可以在自定义层中实现它:

class LocalLinear(nn.Linear):
    def __init__(self, *args, kernel_size=2, **kwargs):
        super().__init__(*args, **kwargs)

        assert self.in_features == kernel_size*self.out_features
        self.mask = torch.block_diag(*[torch.ones(kernel_size,1),]*self.out_features)

def forward(self, input):
    return F.linear(input, self.weight, self.bias)*self.mask

然后像这样定义它:

net = nn.Sequential(nn.Linear(input_size, hidden_size),
                    LocalLinear(hidden_size, output_size))