pytorch batch sparse tensor (构建稀疏batch)

20 篇文章 2 订阅
5 篇文章 0 订阅

直接通过scipy sparse矩阵转化, 这个方法的缺点是不能构建 batch 的sparse tensor

import scipy.sparse as sp
def sparse_mx_to_torch_sparse_tensor(sparse_mx):
       """Convert a scipy sparse matrix to a torch sparse tensor."""
       # sparse_mx = sp.coo_matrix(sparse_mx)
       sparse_mx1 = sparse_mx.tocoo().astype(np.float32)
       indices = torch.from_numpy(
           np.vstack((sparse_mx1.row, sparse_mx1.col)).astype(np.int64))
       values = torch.from_numpy(sparse_mx1.data)
       shape = torch.Size(sparse_mx1.shape)
       return torch.sparse.FloatTensor(indices, values, shape)

直接先建好densor tensor, 再用 to_sparse() 转化就可以得到batch 的sparse

a = torch.randn(10, 2, 3).to_sparse().requires_grad_(True)
b = torch.randn(10, 3, 2)
torch.bmm(a,b)

这个在GNN 中构 建一个batch的邻接矩阵(batch adjcancy matrix)会有用,
https://github.com/pytorch/pytorch/issues/14489
稀疏矩阵乘法可以看: https://blog.csdn.net/weixin_40248634/article/details/107932229

  • 3
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
当涉及到使用稀疏自注意机制的PyTorch代码时,你可以使用`torch_sparse`库中的相关功能。下面是一个简单的示例代码: ```python import torch from torch_scatter import scatter_softmax from torch_sparse import SparseTensor def sparse_self_attention(query, key, value): # 计算注意力得分 attention_scores = torch.matmul(query, key.transpose(-2, -1)) # 使用稀疏张量构建注意力权重矩阵 attention_weights = scatter_softmax(attention_scores, batch_dim=0) sparse_attention_weights = SparseTensor.from_dense(attention_weights) # 应用稀疏注意力权重 output = torch.matmul(sparse_attention_weights, value) return output # 示例输入 query = torch.randn(3, 5, 10) # (batch_size, num_queries, query_dim) key = torch.randn(3, 7, 10) # (batch_size, num_keys, key_dim) value = torch.randn(3, 7, 20) # (batch_size, num_keys, value_dim) # 调用稀疏自注意机制 output = sparse_self_attention(query, key, value) ``` 在上述代码中,我们首先使用`torch.matmul`计算查询向量与键向量之间的注意力得分。然后,我们使用`scatter_softmax`函数计算注意力权重,并使用`SparseTensor.from_dense`将其转换为稀疏张量。最后,我们将稀疏注意力权重应用于值向量,通过稀疏张量的乘法运算实现。这样,我们就得到了使用稀疏自注意机制的输出结果。 请注意,该示例仅展示了如何使用稀疏自注意机制计算注意力。在实际应用中,你可能需要根据具体任务的需求进行相应的修改和适配。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值