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

Question

神经网络密集层中的权重是一个 (n,d) 矩阵，我想强制其中一些权重始终为零。我有另一个 (n,d) 矩阵，它是条目可以为非零的掩码。这个想法是该层不应该是真正密集的，而是缺少一些连接(即等于 0)。
在使用 PyTorch(或 Tensorflow)进行训练时如何实现这一目标？我不希望这些权重在训练时变为非零。
一种方法，如果它不直接支持它，将在每次训练迭代后将所需条目清零。

Answer 1

你可以利用pytorch的sparse数据类型:

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.)