一次零一的节省空间和时间
在之前的博客文章中，我们介绍了稀疏矩阵以及它们如何改善神经网络。
基本的假设是完全密集的层通常过于冗余，可以进行修剪而不会显著损失精度。在某些情况下，稀疏线性层甚至可以提高精度和/或泛化能力。
主要问题是目前可用的支持稀疏代数计算的代码严重缺乏效率。我们还在等待官方的PyTorch支持。
这就是为什么我们在今年夏天耐心耗尽并花费了一些时间来解决这个“空白”的原因。今天，我们很高兴地发布了扩展包pytorch_block_sparse。
单独使用，或者与蒸馏和量化等其他方法结合使用，这个库使得神经网络在生产中既更小又更快，这对于Hugging Face来说是至关重要的，以便让任何人以较低的成本使用神经网络，并提高最终用户的体验。
用法
提供的BlockSparseLinear模块是torch.nn.Linear的直接替代品，使用它在你的模型中非常简单：
# from torch.nn import Linear
from pytorch_block_sparse import BlockSparseLinear
...
# self.fc = nn.Linear(1024, 256)
self.fc = BlockSparseLinear(1024, 256, density=0.1)扩展包还提供了BlockSparseModelPatcher，可以在现有模型上进行“即时”修改，示例笔记本中展示了这一点。然后，可以像往常一样训练这样的模型，而无需更改模型源代码。
NVIDIA CUTLASS
这个扩展包是基于Yulhwa Kim的cutlass tilesparse概念验证的。
它使用了基于CUTLASS的块稀疏矩阵乘法的C++ CUDA 模板。
CUTLASS是一个用于实现高性能CUDA内核的CUDA C++模板集合。通过使用CUTLASS，可以在自定义内核上获得接近cuBLAS性能，而无需使用汇编语言代码。
最新版本包括所有的Ampere Tensor Core原语，在有限的精度损失下提供10倍或更高的速度提升。未来版本的pytorch_block_sparse将利用这些原语，因为块稀疏性与Tensor Core的要求完全兼容。
性能
在当前库的阶段，稀疏矩阵的性能大约比其优化的cuBLAS密集对应物慢两倍，我们相信将来可以改进这一点。
这对于PyTorch稀疏矩阵来说是一个巨大的改进：它们目前的实现比密集矩阵慢一个数量级。
但更重要的是，使用稀疏矩阵的性能提升随着稀疏度的增加而增加，所以一个75%稀疏矩阵大约比密集矩阵快2倍。
内存节省更为显著：对于75%稀疏度，内存消耗减少了4倍，正如你所期望的那样。
未来工作
能够高效训练块稀疏线性层只是第一步。目前，稀疏模式在初始化时是固定的，当然，在学习过程中优化它将会带来巨大的改进。
因此，在未来的版本中，您可以期待工具来测量参数的“有用性”，以便能够优化稀疏模式。在块内，使用NVIDIA Ampere 50%稀疏模式可能会带来另一个显著的性能提升，就像升级到更高版本的CUTLASS一样。
所以，请关注近期更多关于稀疏性的好消息！