没有所有连接的神经网络层

Question

神经网络密集层中的权重是一个 (n,d) 矩阵，我想强制其中一些权重始终为零。我有另一个 (n,d) 矩阵，它是其条目可以非零的掩码。这个想法是，该层不应该是真正密集的，而是缺少一些连接（即等于 0）。

在使用 PyTorch（或 Tensorflow）训练时如何实现这一目标？我不希望这些权重在训练时变得非零。

一种方法（如果不直接支持）是在每次训练迭代后将所需的条目清零。

Answer 1

您可以利用 pytorch 的稀疏数据类型：

class SparseLinear(nn.Module):
  def __init__(self, in_features, out_features, sparse_indices):
    super(SparseLinear, self).__init__()
    self.weight = nn.Parameter(data=torch.sparse.FloatTensor(sparse_indices, torch.randn(sparse_indices.shape[1]), [in_features, out_features]), requires_grad=True)
    self.bias = nn.Parameter(data=torch.randn(out_features), requires_grad=True)

  def forward(self, x):
    return torch.sparse.admm(self.bias, self.weight, x, 1., 1.)