pytorch_block_sparse:用于Pytorch的快速块稀疏矩阵 源码

用于Pytorch的快速块稀疏矩阵 这个PyTorch扩展使用块稀疏矩阵而不是密集矩阵来替代torch.nn.Linear。 由于您可以用稀疏矩阵直接替换模型中的线性层,因此可以非常轻松地进行稀疏矩阵实验。 动机 该库的目的是表明稀疏矩阵可用于神经网络,而不是密集网络,而不会显着改变精度。 这是个好消息,因为稀疏矩阵可释放空间和计算方面的节省: 50%的稀疏矩阵将仅使用50%的内存,并且理论上将仅使用50%的计算。在这个库中,与天真的实现相比,我们利用Cutlass来提高CUDA性能。但是,由于基于cuBLAS的torch.nn.Linear的优化特性,该库的当前版本仍然较慢,大约降低了2倍(将来可能会有所改进)。 在该库的当前阶段,稀疏矩阵的性能比其优化的密集对等矩阵的性能大约慢2倍(我们希望在将来对此进行改进)。但是,使用稀疏矩阵的性能提高随稀疏度的增加而增加,因此75%的稀疏矩阵的速