Closed downeykking closed 10 months ago
orz 太猛了,新实现很优雅,学到很多!想请教一下提高效率的原因是因为 torch_sparse
底层封装了更高效的稀疏矩阵算子吗?因为感觉理论上感觉数据组织形式和原实现相似,但是原实现的计算部分可能都是 naive 的 python 运算?
orz 太猛了,新实现很优雅,学到很多!想请教一下提高效率的原因是因为
torch_sparse
底层封装了更高效的稀疏矩阵算子吗?因为感觉理论上感觉数据组织形式和原实现相似,但是原实现的计算部分可能都是 naive 的 python 运算?
在我个人的理解中,原始实现是进行两步骤,分别是message,aggregate,其中aggregate是通过torch_scatter实现的,从source节点聚合信息到target节点,比如这样:
row, col = edge_index
# x_j为聚合后的x_j,按照0-row.max().item()+1顺序排列
x_j = scatter.scatter(x[col], row, dim=0, dim_size=x.size(0), reduce='mean')
所以这时候要显式提供edge_index.size(1)
份x
,(这个数通常还是非常大,尽管作者对torch_scatter也用c++优化了速度,但显存和速度还是相对不足),并在这个基础上做聚合。
而torch_sparse可以利用message_and_aggregate
变成一步运算,并且聚合的时候也用的是稀疏矩阵和稠密矩阵的乘法,
x_j = matmul(adj_t, x, reduce=self.aggr)
作者写的torch_sparse中用c++重写了很多算子,实现了基于GPU的稀疏矩阵乘法的快速前向和后向传递。
参考:https://pytorch-geometric.readthedocs.io/en/latest/advanced/sparse_tensor.html
原来是这样,学到了,感谢解答!!
为ngcf conv层和lightgcn conv层添加了稀疏矩阵支持,性能基本上维持不变,在测试结果上显存占用变为1/5,提速5x。
对general recommender下面的所有方法都进行了适配(如果用到了ngcf conv或lightgcn conv),所以基于这几个backbone的模型都可以实现提速。为了更好兼容现有代码逻辑,需要手动设置参数才会开启加速。
assert config["enable_sparse"] in [True, False, None]
具体使用方法:
序列推荐其实也有卷积层,也可以进行加速,但是我之前做序列推荐做的比较少,我有时间再研究一下,当作一个todo计划 :)
主要的测试结果: 对于NGCF,15个epoch下,ml-1m数据集 Sparse:占用显存0.18G,训练总时间33.85s Origin:占用显存1.01G,训练总时间155.79s
对于LightGCN,15个epoch下,ml-1m数据集 Sparse:占用显存0.17G,训练总时间23.99s Origin:占用显存1.02G,训练总时间102.95s
对于SGL,15个epoch下,ml-1m数据集 Sparse:占用显存0.60G,训练总时间70.84s Origin:占用显存1.09G,训练总时间382.86s