调用torch.sparse.sum与.to_dense时遇CPU后端错误求助

原始代码 #

class RGCN_Layer(nn.Module):
    """ A Relation GCN module operated on documents graphs. """

    def __init__(self, in_dim, mem_dim, num_layers, relation_cnt=8):
        super().__init__()
        self.layers = num_layers
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.mem_dim = mem_dim
        self.relation_cnt = relation_cnt
        self.in_dim = in_dim

        self.dropout = 0.2
        # self.in_drop = nn.Dropout(self.dropout)
        self.gcn_drop = nn.Dropout(self.dropout)

        # gcn layer
        self.W_0 = nn.ModuleList()
        self.W_r = nn.ModuleList()
        for i in range(relation_cnt):
            self.W_r.append(nn.ModuleList())

        for layer in range(self.layers):
            input_dim = self.in_dim if layer == 0 else self.mem_dim
            self.W_0.append(nn.Linear(input_dim, self.mem_dim).to(self.device))
            for W in self.W_r:
                W.append(nn.Linear(input_dim, self.mem_dim).to(self.device))


    def forward(self, nodes, adj):
        """
        
        :param nodes:  batch_size * num_event * num_event
        :param adj:  batch_size * 8 * num_event * num_event
        :return:
        """
        gcn_inputs = nodes

        maskss = []
        denomss = []
        for batch in range(adj.shape[0]):
            masks = []
            denoms = []
            for i in range(self.relation_cnt):
                denom = torch.sparse.sum(adj[batch, i], dim=1).to_dense()
                t_g = denom + torch.sparse.sum(adj[batch, i], dim=0).to_dense()
                mask = t_g.eq(0)
                denoms.append(denom.unsqueeze(1))
                masks.append(mask)
            denoms = torch.sum(torch.stack(denoms), 0)
            denoms = denoms + 1
            masks = sum(masks)
            maskss.append(masks)
            denomss.append(denoms)
        denomss = torch.stack(denomss) # 40 * 61 * 1

        # sparse rgcn layer
        for l in range(self.layers):
            gAxWs = []
            for j in range(self.relation_cnt):
                gAxW = []

                bxW = self.W_r[j][l](gcn_inputs)
                for batch in range(adj.shape[0]):
                    xW = bxW[batch]  # 255 * 25
                    AxW = torch.sparse.mm(adj[batch][j], xW)  # 255, 25
                    gAxW.append(AxW)
                gAxW = torch.stack(gAxW)
                gAxWs.append(gAxW)
            gAxWs = torch.stack(gAxWs, dim=1)
            gAxWs = F.relu((torch.sum(gAxWs, 1) + self.W_0[l](gcn_inputs)) / denomss)  # self loop
            gcn_inputs = self.gcn_drop(gAxWs) if l < self.layers - 1 else gAxWs

        return gcn_inputs, maskss

错误信息 #

执行到denom = torch.sparse.sum(adj[batch, i], dim=1).to_dense()时触发：

RuntimeError: Could not run 'aten::to_dense' with arguments from the 'CPU' backend. 'aten::to_dense' is only available for these backends: [MkldnnCPU, SparseCPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].

解决方案 #

核心原因 #

报错本质是adj[batch, i]是普通密集张量，而非PyTorch的稀疏张量（torch.sparse.Tensor）。torch.sparse.sum虽能处理密集张量，但返回结果仍是密集张量，此时调用to_dense()属于多余操作——该方法仅适用于稀疏张量。

修复步骤 #

情况1：输入的adj本身就是密集张量

直接替换torch.sparse.sum为普通的torch.sum，并去掉to_dense()：

# 替换原代码中的两行
denom = torch.sum(adj[batch, i], dim=1)
t_g = denom + torch.sum(adj[batch, i], dim=0)

情况2：adj应该是稀疏张量

若业务逻辑要求adj为稀疏格式，需先将其转换为torch.sparse.Tensor（构造时需指定indices和values参数），确保adj[batch, i]是稀疏张量类型后，再执行原代码的torch.sparse.sum和to_dense()操作。

注意：后续的torch.sparse.mm也要求输入为稀疏张量，若adj是密集张量，这一步会隐含转换，影响性能，建议统一张量格式。

内容的提问来源于stack exchange，提问作者Zimu Wang