PyTorch实现稀疏神经网络调用nn.SparseLinear报AttributeError如何解决

错误原因 #

torch.nn模块没有内置SparseLinear类，该报错是API引用错误导致。

可落地的修复方案 #

根据你的使用场景二选一即可：

• 场景1：实现稀疏连接的全连接层（固定比例权重为0，仅剩余权重参与运算更新）
  先自定义SparseLinear层，再替换原有代码的nn.SparseLinear即可，自定义层代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseLinear(nn.Module):
    def __init__(self, in_features, out_features, sparsity=0.9, bias=True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.sparsity = sparsity
        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
        self.bias = nn.Parameter(torch.Tensor(out_features)) if bias else None
        # 权重初始化
        nn.init.kaiming_uniform_(self.weight, a=torch.math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / torch.math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)
        # 生成稀疏掩码，sparsity参数为权重置0的比例
        self.register_buffer('mask', torch.bernoulli(torch.full_like(self.weight, 1 - sparsity)))
    
    def forward(self, x):
        sparse_weight = self.weight * self.mask
        return F.linear(x, sparse_weight, self.bias)

替换后你的模型定义部分代码如下：

class FeedforwardNeuralNetModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, sparsity=0.9):
        super().__init__()
        self.fc1 = SparseLinear(input_dim, hidden_dim, sparsity=sparsity)
        self.sigmoid = nn.Sigmoid()
        self.fc2 = SparseLinear(hidden_dim, output_dim, sparsity=sparsity)

    def forward(self, x):
        out = self.fc1(x)
        out = self.sigmoid(out)
        out = self.fc2(out)
        return out

• 场景2：需要接收COO格式稀疏张量作为层输入
  PyTorch 1.12及以上版本的原生nn.Linear已经支持稀疏张量输入，直接将代码中所有nn.SparseLinear替换为nn.Linear即可正常运行。

内容的提问来源于stack exchange，提问作者user14492868