genbmm:在PyTorch中用于通用矩阵乘法的CUDA内核

Genbmm 该库是PyTorch缺少矩阵乘运算的集合。 它旨在为提供所需的操作员。 该库当前具有两个组件。 它仅支持CUDA操作。 带有梯度的广义矩阵乘法（对数空间，最大值，样本） 带状稀疏矩阵 快速开始 !p ip install - qU git + https : // github . com / harvardnlp / genbmm 广义矩阵乘法 计算矩阵乘以PyTorch中的非标准半环需要创建带有视图的大型中间项。 这是特别糟糕的，因为它们在内部存储用于反向传播。 该库在CUDA中实现了一些替代的矩阵乘法，以避免出现此问题。 import genbmm a = torch . rand ( 10 , 3 , 4 ). cuda (). requires_grad_ ( True ) b = torch . rand ( 10 , 4 , 5 ). cuda ()