PyTorch：计算对称矩阵的上特征向量

Question

我只想找到实数正定对称20x2000x2000 PyTorch张量的主特征向量：

K = torch.random.randn(20, 2000, 2000)
K_sym = K+K.transpose(1,2)/2
eig, ev = torch.symeig(K_sym, eigenvectors=True)
ev = ev[:,1999,:]

ev.shape

> torch.Size[20, 2000]

这种实现非常缓慢，即使在GPU上

有人能提出一种更快的方法来找到第一个特征向量吗

例如，幂迭代可能是一种解决方案。我是根据torch.nn.utils.spectral_norm中的解决方案得出这一结论的：

def TopEigenvector(self, K, n_power_iterations=400, dim=1, eps=1e-12):
        v = torch.ones(K.shape[0],K.shape[1],1).to(self.device)
        with torch.no_grad():
            for _ in range(n_power_iterations):
                v = normalize(torch.bmm(K, v),dim=dim, eps=eps, out=v)
                if(n_power_iterations > 0):
                    v = v.clone(memory_format=torch.contiguous_format)
        return v

为了适应这个问题，我将其更改如下：

    def TopEigenvector(self, K, n_power_iterations=400, dim=1):
        v = torch.ones(K.shape[0],K.shape[1],1).to(self.device)
        for _ in range(n_power_iterations):
            m = torch.bmm(K, v)
            n = torch.norm(m, dim=1).unsqueeze(1)
            v = m/n
        return v

然而，这依赖于可怕的for循环，仍然很慢，但比原始方法快

有没有更好更快的方法可以使用