如何在PyTorch中高效实现非全连接的线性层？

Question

我制作了一个我想要实现的缩小版本的示例图：

因此顶部的两个输入节点仅与顶部的三个输出节点完全连接，同样的设计也适用于底部的两个节点。到目前为止，我已经想出了两种在 PyTorch 中实现这一点的方法，但这两种方法都不是最佳的。

第一个方法是创建一个包含许多较小线性层的 nn.ModuleList，并在前向传递过程中，通过它们迭代输入。对于该图的示例，它看起来像这样：

class Module(nn.Module):
  def __init__(self):
    self.layers = nn.Module([nn.Linear(2, 3) for i in range(2)])
  
  def forward(self, input):
    output = torch.zeros(2, 3)
    for i in range(2):
      output[i, :] = self.layers[i](input.view(2, 2)[i, :])
    return output.flatten()

这样就完成了图中的网络，主要问题是速度非常慢。我认为这是因为 PyTorch 必须顺序处理 for 循环，并且无法并行处理输入张量。

为了“矢量化”模块以便 PyTorch 可以更快地运行它，我有这个实现：

class Module(nn.Module):
  def __init__(self):
    self.layer = nn.Linear(4, 6)
    self.mask = # create mask of ones and zeros to "block" certain layer connections
  
  def forward(self, input):
    prune.custom_from_mask(self.layer, name='weight', mask=self.mask)
    return self.layer(input)

这也完成了图的网络，通过使用权重修剪来确保全连接层中的某些权重始终为零（例如，连接顶部输入节点到底部输出节点的权重始终为零，因此它实际上“断开连接”） ）。该模块比之前的模块快得多，因为没有 for 循环。现在的问题是这个模块占用了更多的内存。这可能是因为，即使大多数层的权重为零，PyTorch 仍然将网络视为就好像它们在那里一样。这个实现本质上保留了比它需要的更多的权重。

有人遇到过这个问题并提出有效的解决方案吗？

Answer 1

如果权重共享没问题，那么一维卷积应该可以解决问题：

class Module(nn.Module):
  def __init__(self):
    self.layers = nn.Conv1d(in_channels=2, out_channels=3, kernel_size=1)
    self._n_splits = 2

  
  def forward(self, input):
    
    B, C = input.shape
    output = self.layers(input.view(B, C//self._n_splits, -1))
    return output.view(B, C)

如果权重共享不行，那么你可以使用组卷积：

self.layers = nn.Conv1d(in_channels=4, out_channels=4, kernel_size=1, stride=1, groups=2)

。不过，我不确定这是否可以实现任意数量的通道分割，你可以查看文档：https://pytorch.org/docs/stable/ generated/torch.nn.Conv1d.html

一维卷积是输入的所有通道上的全连接层。组卷积会将通道分为组并对它们执行单独的卷积操作（这就是您想要的）。

实现将类似于：

class Module(nn.Module):
  def __init__(self):
    self.layers = nn.Conv1d(in_channels=2, out_channels=4, kernel_size=1, groups=2)

  
  def forward(self, input):
    
    B, C = input.shape
    output = self.layers(input.unsqueeze(-1))
    return output.squeeze()

编辑：

如果您需要奇数个输出通道，您可以组合两组卷积。

class Module(nn.Module):
  def __init__(self):
    self.layers = nn.Sequence(
         nn.Conv1d(in_channels=2, out_channels=4, kernel_size=1, groups=2),
         nn.Conv1d(in_channels=4, out_channels=3, kernel_size=1, groups=3))


  def forward(self, input):
    
    B, C = input.shape
    output = self.layers(input.unsqueeze(-1))
    return output.squeeze()

这将根据图中的要求有效地定义输入通道，并允许您使用任意数量的输出通道。请注意，如果第二个卷积具有

groups=1

，您将允许混合通道，并且将有效地使第一组卷积层变得无用。

从理论角度来看，这两个卷积之间不需要激活函数。我们将它们组合成线性物质。然而，添加激活函数可能会提高性能。

Answer 2

杰克，抱歉这么晚回复。 如果您的输入是 B（等于 2xb 的偶数）并且您的输出是 C，您可以将其重塑为 2 x b，然后简单地将线性层 (b,C) 应用到最后一个维度吗？