hkhkhkzxk的博客

原创 pytorch自定义用于稀疏矩阵乘法的反向传播函数

一、 整体代码如上，分为两个部分：1、一个继承了torch.autograd.function的类，类中有两个静态函数,forward和backward ； 2 、一个实现了上述类的spmm类，用于调用并且返回梯度结果，两个部分结合使用的效果相当于一个函数。这个函数的作用是实现稀疏矩阵乘法的前向传播和反向传播。它接收四个输入参数：indices（稀疏矩阵的非零元素的位置信息），values（稀疏矩阵的非零元素的值），shape（矩阵的形状）和b（矩阵乘法的运算数）。二、在前向传播中，它首先使用indices