self.loop_times = loop_times self.fuse_alpha = fuse_alpha def k_hop(self, A): # A: N, C, V, V N, C, V, _ = A.shape # A0: 1, 1, V, V A0 = torch.eye(V, dtype=A.dtype).to(A.device).unsqueeze(0).unsqueeze(0) * self.fuse_alpha A_power = torch.eye(V, dtype=A.dtype).to(A.device).unsqueeze(0).unsqueeze(0) for i in range(1, self.loop_times + 1): A_power = torch.einsum('ncuv,ncvw->ncuw', A, A_power) A0 = A_power * (self.fuse_alpha * (1 - self.fuse_alpha) ** i) + A0 return A0将此代码换一种写法

这段代码的功能是实现了一个 k-hop 图卷积操作，可以将节点之间的连接关系进行卷积操作，得到不同 k-hop 的邻接矩阵。

以下是另一种可能的实现方式：

class GraphConv(nn.Module):
    def __init__(self, loop_times, fuse_alpha):
        super(GraphConv, self).__init__()
        self.loop_times = loop_times
        self.fuse_alpha = fuse_alpha

    def k_hop(self, A):
        # A: N, C, V, V
        N, C, V, _ = A.shape

        # A0: 1, 1, V, V
        A0 = torch.eye(V, dtype=A.dtype, device=A.device).unsqueeze(0).unsqueeze(0) * self.fuse_alpha
        A_power = A.clone()

        for i in range(1, self.loop_times + 1):
            A_power = torch.matmul(A_power, A)
            A0 += A_power * (self.fuse_alpha * (1 - self.fuse_alpha) ** i)

        return A0

这里使用了 `nn.Module` 的子类来实现。在 `__init__` 函数中定义了 `loop_times` 和 `fuse_alpha` 两个参数，表示 k-hop 的次数和权重系数。在 `k_hop` 函数中，首先获取输入张量 A 的形状，并构造一个大小为 1 × 1 × V × V 的张量 A0，用于存储所有 k-hop 邻接矩阵的加权和。然后，通过 `torch.eye` 函数创建一个大小为 V × V 的单位矩阵，并将其扩展为 1 × 1 × V × V 的形状，与权重系数相乘，得到大小为 1 × 1 × V × V 的初始邻接矩阵 A0。

在循环中，计算当前 k-hop 的邻接矩阵 A_power，然后与 A0 进行加权相加，得到更新后的 A0。最后返回 A0，即为从 0-hop 到 loop_times-hop 的邻接矩阵的加权和。